🧬 biology

Characterizing optimal hierarchical policy inference on graphs via non-equilibrium thermodynamics

This paper introduces a formalism based on non-equilibrium thermodynamics to derive optimal state-space hierarchies for discrete Markov decision processes on graphs, framing the resulting policy inference as a hierarchical gradient flow between prior and optimal trajectory densities.

Original authors: Daniel McNamee

Published 2026-06-04

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Daniel McNamee

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ⚕️ This is an AI-generated explanation of a preprint that has not been peer-reviewed. It is not medical advice. Do not make health decisions based on this content. Read full disclaimer

Imagine you are trying to find the best route through a giant, complex maze. You have a map (the "prior" policy), but it's just a guess. You know there are rewards at the end of certain paths, but you don't know exactly which way to turn to get there most efficiently.

This paper proposes a new way to understand how a smart agent (like a human or a robot) figures out the best path. Instead of just calculating one step at a time, it looks at the entire journey as a flowing river of possibilities.

Here is the breakdown using simple analogies:

1. The "River of Possibilities" (The Setup)

Think of every possible path you could take through the maze as a tiny particle floating in a river.

The Prior Policy: At the start, these particles are spread out randomly, representing your initial guesses or habits.
The Reward: Imagine the maze has "gravity" pulling everything toward the exit (the reward). The better the path, the stronger the pull.
The Goal: We want all these particles to eventually settle into the single, perfect path that gets you to the reward with the least amount of wasted effort.

2. The Physics of Thinking (Non-Equilibrium Thermodynamics)

The author uses a concept from physics called thermodynamics to describe how thinking works.

Imagine the particles are hot gas molecules. They are jiggling around randomly.
The "reward" acts like a cooling system. As the particles move, they naturally drift toward the "coolest" (most rewarding) spots.
The paper suggests that the process of planning is just watching this gas cool down and settle into the perfect shape. It's not a sudden jump; it's a smooth flow from a messy guess to a perfect solution.

3. The "Flow" of Decisions (Policy Inference)

The paper introduces a mathematical rule (the Fokker-Planck equation) that describes how this flow happens.

Think of it like water flowing down a hill. The water naturally finds the steepest, fastest path to the bottom.
In our maze, the "water" is your decision-making process. It flows from your initial confusion toward the optimal path.
Crucially, this flow happens across all possible paths at once, not just one. It considers how every single step connects to every other step, creating a "hierarchy" of importance.

4. Finding the "Bottlenecks" (The Hierarchy)

This is the most important part of the discovery. As the "water" flows, it speeds up at certain points and slows down at others.

The Bottleneck: Imagine a narrow bridge connecting two large rooms in the maze. Almost everyone has to cross this bridge to get to the other side.
The paper shows that this mathematical flow naturally highlights these bottlenecks. These are the most important states in the maze.
Why it matters: If you are trying to solve the maze, you should focus your attention on these bottlenecks first. They are the "keys" to the whole structure. The paper claims that by following this flow, an agent automatically learns to prioritize these critical junctions, creating a mental hierarchy of the maze.

5. The Experiment (The Regular Graph)

To test this, the author used a specific type of maze (a regular graph) that looks very uniform and boring—every spot looks the same, with no obvious landmarks.

The Human Test: In previous studies, humans were asked to find the shortest path in this maze. Even though the maze looked uniform, humans intuitively identified the "bottleneck" bridge as the most important spot.
The Computer Test: The author ran their "flow" math on the same maze. The math identified the exact same bottleneck as the most important spot.
The Result: When the computer used this "hierarchical" order to plan (checking the bottlenecks first), it solved the maze much faster and with less confusion than if it had checked random spots. It was like having a GPS that told you, "Don't worry about the side streets; focus on the bridge."

Summary

The paper argues that optimal planning is like a physical flow. By treating decision-making as a fluid moving toward a reward, we can mathematically prove that the best way to solve a problem is to identify the "bottlenecks" or critical junctions first. This creates a natural hierarchy, allowing a brain or a computer to ignore the noise and focus on the most important parts of the map, just like a human does intuitively.

Technical Summary: Characterizing Optimal Hierarchical Policy Inference on Graphs via Non-Equilibrium Thermodynamics

Problem Statement
Hierarchies are recognized as fundamental to both stochastic optimal control and biological control systems, offering computational advantages and potentially serving as core principles for sensorimotor and cognitive control. However, a theoretically justified method for constructing state-space hierarchies across all spatial resolutions, and specifically characterizing their evolution through a policy inference process, has remained elusive. The paper addresses the challenge of deriving normative representations of discrete Markov Decision Processes (MDPs) that reveal these hierarchical structures without relying on ad-hoc constructions.

Methodology
The authors propose a formalism that interprets policy inference in discrete MDPs through the lens of non-equilibrium thermodynamics.

MDP Formulation as Free Energy Minimization:
The standard MDP objective of maximizing cumulative reward is reformulated as a free energy functional $J[\pi]$ . By embedding the MDP in a linearly solvable framework (KL Control), the objective incorporates a trade-off between expected reward and policy complexity (entropy). The optimal policy $\pi^*$ is expressed via a Gibbs-Boltzmann distribution, balancing a prior policy $\pi_{prior}$ against the reward signal.
Thermodynamic Interpretation:
Planning is conceptualized as a physical process where "particles" representing state trajectories evolve from an initial density (the prior policy) toward a thermal equilibrium (the optimal policy). This evolution is governed by a potential $\Psi(x) = -R(x)$ . Instead of simulating stochastic particle dynamics, the paper analyzes the evolution of the trajectory density $\pi_\tau$ over time $\tau$ .
Fokker-Planck Dynamics:
The evolution of the policy density is described by the Fokker-Planck (FP) equation. The authors establish that the time evolution of $\pi_\tau$ according to the FP equation corresponds to a steepest descent on the free energy functional $J[\pi]$ with respect to the Wasserstein metric. This implies that the FP equation describes the dynamics of an optimal policy inference algorithm transforming the prior into the optimal policy.
Discretization and State-Space Projection:
To apply this to discrete MDPs, the continuous FP equation is discretized using a Markov chain approximation of the underlying Langevin dynamics. This yields a discrete FP equation (Eqn. 5) driven by a generalized potential $\bar{\Psi}(x) = -R(x) + \log \pi(x)$ .
Crucially, the authors "pull back" the trajectory-level flux to the state level. By summing the total flux at each state, they derive a dynamic state hierarchy. This is expressed geometrically using a discrete Laplacian operator $\Delta_D$ defined by the fundamental matrix of the policy's Markov chain:
$\dot{\pi}(x) = \Delta_D J_\pi(x)$
Here, $J_\pi(x)$ represents the local contribution to the free energy (planning complexity) at state $x$ .

Key Contributions and Results

Normative Hierarchical Representation: The paper introduces a formal derivation of state-space hierarchies as a byproduct of optimal policy inference. The hierarchy is not pre-defined but emerges from the dynamics of the policy flow.
Identification of Bottlenecks: In an application to a regular graph (a common benchmark with homogeneous local structure but global bottlenecks), the proposed metric $\Delta_D J_\pi(x)$ successfully identifies bottleneck states. The magnitude of the policy flux at these states corresponds to their rank in the optimal hierarchy.
Information-Theoretic Validation: The authors demonstrate that the hierarchical ordering derived from the flux metric aligns with human behavioral data regarding bottleneck identification. Furthermore, they show that an "oracle planner" utilizing this hierarchical sampling order reduces the entropy of candidate trajectories (i.e., gains information) more efficiently than random sampling. The highest-ranked states in the hierarchy provide the most information regarding the shortest path.

Significance and Claims
The paper claims to establish a principle of optimality for policy inference processes: the Fokker-Planck equation describes the expected time evolution of processes that take the steepest descent on the free energy objective. By operating in trajectory space, this approach integrates state-space structure at all scales into the inference process.

The authors position this work as a preliminary study offering a "normative perspective" on optimal hierarchical policy inference. They suggest that the framework provides a theoretical link between optimal control and hierarchical representations in human cognition. The significance lies in the potential to explain aspects of behavioral and neural dynamics in natural planning and problem-solving, though the authors modestly note that embedding these principles into practical algorithms remains an open question. The work does not claim to solve all hierarchical learning problems but rather provides a formal characterization of how optimal hierarchies should theoretically evolve.