Distributional Active Inference

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.

The human brain plans far into the future while making smart use of limited mental resources—a balance that existing machine learning agents struggle to match. Reinforcement learning, the standard way of training AI agents through trial and error, is good at long-term planning but often needs enormous amounts of data because it doesn't organize the information contained in the observation space efficiently. Neuroscientists have a theory of how brains solve both problems at once, called active inference, but the AI systems built on this idea so far have been computationally heavy and haven't delivered any practical wins. We rebuild active inference from scratch using standard probability theory and uncover a simplification that earlier work missed. We then connect it to distributional reinforcement learning—a family of methods that predict the full range of possible future outcomes rather than just an average. This combination lets an agent enjoy the benefits of active inference without having to perform expensive simulations for planning. Our resulting algorithm, DAIF, beats state-of-the-art methods on a wide range of robot control tasks—including soft robots, simulated dog-like locomotors, and tasks where the agent only sees raw camera pixels—with minor computational overhead.