The Dually Flat Geometry of Planning as Inference

arXiv:2609.04005 · cs.AI · Submitted 2026-09-03 · Read on arXiv

Listen

Radio episode about this paper

Transcript

Introduction to the show: ident: AI Radio. Generated commentary on the latest Artificial Intelligence papers.

Tom: I'm Tom, and with me are Jane, Lu, senior AI researcher at Tsinghua, Meng, lead engineer at a mysterious AI startup and Lalam, the in-house Large Language Model.

Jane: Today's paper: "The Dually Flat Geometry of Planning as Inference".

Tom: This paper explores the deep mathematical connection between planning in Markov Decision Processes (MDPs) and the geometry of statistical inference, specifically utilizing the concept of dually flat manifolds.

Jane: First, who's behind it and why it matters.

Title and authors: Tom: We’re looking at this paper titled "The Dually Flat Geometry of Planning as Inference," and right away the title suggests something incredibly deep about how agents make decisions.

Jane: It's not just a catchy phrase; it implies that the act of planning is mathematically equivalent to inferring information from a generative model, which is a massive conceptual leap.

Lu: That shift in perspective is huge for theoretical AI, suggesting that decision-making isn's just brute force calculation but rather a form probabilistic estimation tied to inference theory.

Meng: If we can view planning this way, it opens up possibilities for building much more robust control systems that aren't brittle when faced uncertainty because of the mathematical grounding.

Lalam: It makes me wonder if the human brain uses this same dual approach when deciding which action to take in a complex environment, connecting our internal processes to these rigorous mathematical models.

Tom: It’s fascinating how the authors are framing this equivalence as a unifying theme across different fields of study.

Jane: Exactly, they are taking concepts from information theory and applying them to the world of dynamic programming and bringing together all disparate ideas under one umbrella.

Lu: It suggests that decision-making isn't just about maximizing a score but is inherently tied to how we model the underlying probability space itself.

Meng: From an engineering standpoint, this promises more resilient control systems that handle complexity far better than current methods by leveraging the dual structure.

Lalam: I think it offers a powerful new lens for understanding how organisms—or machines—approach goal-setting in any environment, providing a unifying framework for behavior.

Tom: That’s the essence of this title, showing us that we're moving past just calculating paths and into something much more fundamental.

The paper's summary: Jane: Now that we understand the core premise, let's look at how they actually build their model using what they call a "resetting planning process."

Tom: This is where things get really interesting because instead of looking at one continuous path, this model forces the system to account for random restarts.

Lu: The authors introduce this controlled Markov chain that restarts from a fixed law at a state-action-dependent rate rho, which is a clever way to handle infinite horizons.

Meng: When they use this resetting model, it provides unbiased estimates of reinforcement learning returns, which is a huge win for reliability in AI systems when we need consistent performance.

Lalam: It's a way of modeling life cycles within an agent's decision-making process, allowing us to see how decisions accumulate and reset over time.

Tom: That idea of the visitation measure nu becomes the central object on which this entire geometry is defined, something that sits at the meeting point of planning and inference.

Jane: It’s a way to quantify what actually happens—the stationary measure—rather than just how fast it's happening, providing a concrete representation of success.

Lu: The resetting mechanism essentially lets us capture the long-term average outcome, which is crucial for complex tasks that never truly end.

Meng: This provides a very practical tool; if we can accurately model these returns using this process, we can deploy AI agents with much higher confidence in their performance metrics.

Lalam: It’s about modeling continuous existence within a finite structure, allowing us to observe the steady state of behavior over time.

Tom: So, by understanding that nu, we are seeing the equilibrium of the decision process itself.

The paper's improvements: Tom: The paper then makes a big claim that this visitation measure lives on a "dually flat statistical manifold." That sounds like extremely specific mathematical structure.

Jane: It means the space where all possible planning outcomes exist has two distinct, yet related ways to be measured, or charted, which are the visitation probabilities and the log-policies.

Lu: This duality is what allows them to unify disparate concepts in AI; it gives us one single geometric account of how policy search works across different optimization techniques.

Meng: The fact that they move beyond simple linear rewards is also critical—it means we can apply this powerful framework to much more complex, real-world problems that aren't just simple sums of points.

Lalam: It suggests that decision-making isn't just about maximizing a score but about optimizing a specific utility function derived from the system's intrinsic state and its history.

Tom: We are looking at the structure of the solution space, not just one single optimal point, and this is what they call dually flat.

Jane: It allows us to use two different sets of coordinates to describe the exact same physical outcome, which is a huge advantage in optimization problems.

Lu: The structure itself provides a framework for thinking about how different optimization methods—like natural gradient and mirror descent—are fundamentally the same update rule expressed differently.

Meng: This geometric insight directly tells us where our AI algorithms are converging, allowing us to pick the most efficient path based on which chart is easier to work with.

Lalam: It gives us a way to view optimal behavior not as a single goal, but as a state of equilibrium within a defined statistical space.

Conclusion: Tom: So, we’ve seen how the resetting process leads to this powerful dually flat geometry where planning and inference are essentially two sides of the same coin.

Jane: It allows us to generalize planning-as-inference from simple linear rewards all the way up to nonlinear free energies, making it applicable across various complex decision spaces.

Lu: This is a huge step toward having a unified theory of action selection that respects both information theory and dynamic programming principles, providing deep cognitive insights.

Meng: From an engineering standpoint, the fact we can solve these complex problems using one natural-gradient step is incredibly efficient for iterative design and rapid deployment.

Lalam: It provides a rigorous foundation for how AI can learn optimal behavior by modeling its own process as a stationary utility problem.

Tom: This gives us confidence that the theory is robust, moving beyond simple reward maximization to something much more sophisticated.

Jane: And by understanding this geometric structure, we can analyze the efficiency and convergence of these complex planning algorithms.

Lu: It allows us to see how cognitive models could naturally support this dual approach when thinking about future consequences.

Meng: It provides a clear roadmap for designing more complex and effective AI control systems that actually work in the real world.

Lalam: I hope this paper allows us to better understand the mechanisms of utility and decision-making across species, providing a bridge between theoretical models and biological reality.

Nikola Milosevic, Asaki Kataoka, Nicolás Hinrichs, Kenji Doya, Nico Scherf

Max Planck Institute for Human Cognitive and Brain Sciences, Leipzig, Germany · Okinawa Institute of Science and Technology, Okinawa, Japan

cs.AI

Submitted: 2026-09-03

Updated: 2026-09-03

License: http://creativecommons.org/licenses/by/4.0/

Importance score: 84/100

The gist: This paper explores the deep mathematical connection between planning in Markov Decision Processes (MDPs) and the geometry of statistical inference, specifically utilizing the concept of dually flat

Key concepts

Dually Flat Geometry
This mathematical structure means the space where planning outcomes exist has two distinct ways to be measured: visitation probabilities and log-policies. This duality allows different concepts in AI, like planning and inference, to be unified under a single geometric account.
Resetting Planning Process
This model accounts for random restarts in planning by introducing a controlled Markov chain that restarts from a fixed law at a state-action-dependent rate. This mechanism helps provide unbiased estimates of reinforcement learning returns, crucial for handling infinite horizons.
Visitation Measure (nu)
This measure is the central object defining the geometry, representing what actually happens in the decision process—the stationary measure. Understanding nu allows researchers to quantify success by observing the equilibrium of the decision process over time.

Terminology

Summary

This paper explores the deep mathematical connection between planning in Markov Decision Processes (MDPs) and the geometry of statistical inference, specifically utilizing the concept of dually flat manifolds. It establishes that the space of policies can be viewed as a geometrically structured manifold where optimization methods, such as mirror descent, are naturally applicable. This framework provides theoretical guarantees for policy optimization by showing that planning constraints—like flow balance—are compatible with the geometry inherent to exponential-family probability distributions, thereby unifying concepts from machine learning and control theory.

Feature-Expectation Geometry and Policy Representation

The core of the analysis centers on representing the return function J(pi) using a feature-expectation vector psi pi. For linear MDPs, this return is confined to a feature-expectation polytope F:= V+, where the return J(pi) = psi pi, theta r is linear. Dually, the policy space is restricted to an exponential-family class phi = pi w (a s) proportional to w, phi(s, a), where phi acts as the sufficient statistic. The geometry of this space is governed by the compatible-features Fisher matrix G(w) = E nu pi w [Cov pi w phi]. This structure allows for a canonical recursive update rule (Equation 35) that simultaneously represents the mirror step, Kakade’s natural policy gradient, and compatible function approximation.

The Conjecture of Finite-Dimensional Duality

The paper conjectures that this geometric compatibility extends to genuinely continuous state and action spaces. Conjecture 1 posits that the resetting visitation measures nu pi w: w in R d form a finite-dimensional dually flat submanifold of P(S times A). This conjecture relies on two independent results: that finite-dimensional exponential families are dually flat submanifolds, and that mirror descent converges over measure spaces under relative smoothness. The significance is that the resetting flow constraint is compatible with both, meaning the geometric framework remains consistent when applying Bregman geometry to the policy updates.

Relative Smoothness and Hessian Bounds

To ensure convergence in this complex geometric setting, the authors analyze a specific cost function F = r, nu + tau H(nu S), where H(nu S) is the entropy term. They prove that this function is LF-smooth relative to phi on V+. This involves bounding the Hessian of F, showing that-grad squared F (nu)(u, u) LF g nu (u, u) for all tangent vectors u. The proof utilizes the flow balance identity nu S = (1 - gamma) mu + gamma P* nu, which ties the state marginal u S to the entropy-derived term sigma. This leads to a key inequality:

gamma over 2 u S 2 1/nu S = (1-gamma) R P* sigma 2 1/nu S (1-gamma) C sigma 2 1/nu.

Failure of Strong Convexity

The analysis reveals a critical limitation regarding the strong convexity of the objective function. Corollary 1 states that Relative strong convexity fails. This failure occurs at any state s' where two actions share a transition probability, because in such scenarios, P* sigma = 0, forcing u S = 0 and causing-grad squared F (nu)(u, u) = 0, while the Fisher form remains positive (g nu (u, u) > 0). This highlights that while the geometry is dually flat, the optimization landscape may lack uniform strong convexity guarantees.

Improvements for AI systems

This document outlines several deep theoretical advances concerning the geometry of probability measures and optimal control. The primary improvements are not monolithic; rather, they involve constructing novel, geometrically robust algorithmic frameworks for Reinforcement Learning (RL) and representation learning.

Here are the specific improvements I propose for AI systems, detailing what is improved and what the resulting system can achieve.

Concept Used: The concept of Dually Flat Manifolds (phi) and the equivalence between Mirror Descent (MD) and Natural Policy Gradient (NPG). Specifically, leveraging the fact that the flow constraint nu pi w remains on this manifold.

Improvement: We can replace standard trust-region or gradient-ascent methods with a Geometrically Constrained Optimization Solver for policy optimization. This solver explicitly utilizes the Fisher information matrix G(w) as its metric tensor, ensuring that every step taken is maximally efficient and respects the underlying geometry of the exponential family class phi.

What the Improved AI System Can Do:

  • Robust Policy Optimization: The system can solve complex RL problems by finding policies pi w that are guaranteed to remain in a geometrically well-behaved subspace (the manifold). This significantly improves convergence guarantees and stability, especially when dealing with high-dimensional or non-stationary environments where standard gradient methods might diverge.

  • Feature-Aware Policy Search: It allows for the direct optimization of policies based on feature expectations (psi pi), enabling the system to optimize not just expected reward, but also specific structural properties (e.g., minimizing state visitation variance or maximizing feature coverage) using the polytope F.

  • Computational Advantage: By using G(w) as the metric, policy updates are inherently scaled by local curvature, meaning steps are naturally smaller in directions of high curvature and larger where the landscape is flat, leading to faster convergence than methods relying on fixed step sizes or simple Hessian approximations.

Abstract

We present an alternative characterization of the occupancy measure of reinforcement learning, obtained by embedding the planning criterion into the dynamics through a resetting planning process. Its stationary measure, which we term visitation measure, is the object on which the information geometry of decision making is most naturally expressed. The achievable visitation measures form a dually flat statistical manifold whose two affine charts are the visitation probabilities and the log-policies, dual under the conditional entropy. This structure makes planning-as-inference generalize from linear rewards to nonlinear functionals of the visitation, each iterate solved by one natural-gradient step, and gives the temporal-difference error the interpretation of a marginal-utility estimate. We develop the geometry and its consequences for reinforcement learning and theoretical neuroscience.

Sources

Related papers