Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Distributional Active Inference

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

TL;DR

This paper addresses how to combine Active Inference with distributional reinforcement learning to enable far-sighted planning under limited computation. It develops a three-part framework: a rigorous AIF formulation, a push-forward RL perspective that links model-based and model-free views, and a practical, model-free algorithm called Distributional Active Inference (DAIF) that uses a latent encoder and quantile regression of latent returns. The authors prove contraction properties for the distributional Bellman operator under push-forward mappings and show how latent encodings can speed up convergence when dynamics compress onto a latent manifold. Empirically, DAIF improves performance across tabular and continuous control tasks, including RiverSwim, EvoGym, and DeepMind Control tasks, while incurring modest additional compute relative to strong distributional baselines. The work offers a scalable route to harness AIF’s information-foraging and exploration benefits in practical RL settings without learning a full world model.

Abstract

Optimal control of complex environments with robotic systems faces two complementary and intertwined challenges: efficient organization of sensory state information and far-sighted action planning. Because the reinforcement learning framework addresses only the latter, it tends to deliver sample-inefficient solutions. Active inference is the state-of-the-art process theory that explains how biological brains handle this dual problem. However, its applications to artificial intelligence have thus far been limited to extensions of existing model-based approaches. We present a formal abstraction of reinforcement learning algorithms that spans model-based, distributional, and model-free approaches. This abstraction seamlessly integrates active inference into the distributional reinforcement learning framework, making its performance advantages accessible without transition dynamics modeling.

Distributional Active Inference

TL;DR

This paper addresses how to combine Active Inference with distributional reinforcement learning to enable far-sighted planning under limited computation. It develops a three-part framework: a rigorous AIF formulation, a push-forward RL perspective that links model-based and model-free views, and a practical, model-free algorithm called Distributional Active Inference (DAIF) that uses a latent encoder and quantile regression of latent returns. The authors prove contraction properties for the distributional Bellman operator under push-forward mappings and show how latent encodings can speed up convergence when dynamics compress onto a latent manifold. Empirically, DAIF improves performance across tabular and continuous control tasks, including RiverSwim, EvoGym, and DeepMind Control tasks, while incurring modest additional compute relative to strong distributional baselines. The work offers a scalable route to harness AIF’s information-foraging and exploration benefits in practical RL settings without learning a full world model.

Abstract

Optimal control of complex environments with robotic systems faces two complementary and intertwined challenges: efficient organization of sensory state information and far-sighted action planning. Because the reinforcement learning framework addresses only the latter, it tends to deliver sample-inefficient solutions. Active inference is the state-of-the-art process theory that explains how biological brains handle this dual problem. However, its applications to artificial intelligence have thus far been limited to extensions of existing model-based approaches. We present a formal abstraction of reinforcement learning algorithms that spans model-based, distributional, and model-free approaches. This abstraction seamlessly integrates active inference into the distributional reinforcement learning framework, making its performance advantages accessible without transition dynamics modeling.
Paper Structure (37 sections, 5 theorems, 66 equations, 8 figures, 6 tables, 7 algorithms)

This paper contains 37 sections, 5 theorems, 66 equations, 8 figures, 6 tables, 7 algorithms.

Table of Contents

  1. Introduction
  2. A Rigorous Formulation of Active Inference
  3. Push-forward Reinforcement Learning
  4. Active Inference with Push-Forward RL
  5. The Prior Art of RL with Active Inference
  6. Experiments
  7. Takeaways and Open Questions
  8. Proofs and derivations
  9. Proof of \ref{['lem:push-forward-contraction']}
  10. Proof of Lemma \ref{['lem:composite_lipschitz']}
  11. Proof of \ref{['thm:ppi-convergence']}
  12. Factorization of the intervened world model
  13. Derivation of the active inference objective
  14. Derivation of the upper bound to the push-forward RL objective
  15. Experiment Details
  16. ...and 22 more sections

Key Result

Lemma 3.1

Let $P_\pi,\bar{P}_\pi,P_\pi^*$ be Markov kernels induced by a fixed policy $\pi$ and $\mathbf F$ be a push-forward process functional, then the distributional Bellman operator ${\bf T}_{P_*}^\pi$ is a contraction with respect to $\bar{\mathcal{W}}_p$

Figures (8)

  • Figure 1: Evaluation curves for three representative environments, one per suite, where DAIF clearly improves the state of the art. These are relatively harder problems of the related suite due to either complex dynamics or large state or action dimensionality, where the abstraction of the return distribution is beneficial. DAIF performs comparably to the state of the art when it does not improve. For the learning curves of the remaining environments, see \ref{['appsec:cont_control']}.
  • Figure 2: Horizontal axis: Distance from the initial state to the most desired state. Vertical axis: Frequency of the visitation of the most desired state. DAIF matches the plain distributional RL performance when transition dynamics cannot be represented more efficiently in a latent space (left panel). DAIF outperforms both distributional and model-based counterparts when a latent manifold drives the dynamics at a degree increasing with the difficulty of the problem (right panel). For the learning curves of individual configurations, see \ref{['fig:tabular_results']} in the appendix.
  • Figure 3: Observation-to-latent state mappings in the Latent RiverSwim environment across multiple horizons. The $(i,j)$ coordinates represent indices for each observation in $\mathcal{X}$, while cell values and colours indicate the corresponding latent state index from $\mathcal{S}$. For instance, at Horizon = 4, the cell $(i=3, j=1)$ has value $2$, signifying that observation $(3,1)$ is mapped to latent state $2$ by the encoding process. This visualization demonstrates that when the transitions and rewards originate from a lower-dimensional manifold, the encoder naturally forms equivalence classes among observations through state abstraction. This abstraction reduces the complexity of the mapping, inducing a more favourable Lipschitz constant for the decoder.
  • Figure 4: Top: RiverSwim. Bottom: Latent RiverSwim. Solid lines indicate the mean over $50$ seeds, and shaded regions represent the standard error.
  • Figure 5: Learning curves for EvoGym environments.
  • ...and 3 more figures

Theorems & Definitions (10)

  • Lemma 3.1
  • Definition 3.2
  • Definition 3.3
  • Lemma 3.4
  • Theorem 3.5
  • Lemma 1.1
  • proof
  • Lemma 1.2
  • proof
  • proof : Proof of \ref{['thm:ppi-convergence']}.