PaperlayerView exact revision on arXiv

Source-linked AI summary

The Dually Flat Geometry of Planning as Inference

Nikola Milosevic, Asaki Kataoka, Nicolas Hinrichs, Kenji Doya, Nico Scherf

arXiv:2609.04005v1cs.AI

TL;DR

The paper asks how planning-as-inference can share a principled variable and geometry with reinforcement-learning decision making. It introduces a resetting planning process, identifies its stationary visitation measure as that variable, and shows that the resulting dually flat manifold supports nonlinear planning updates and a marginal-utility interpretation of temporal-difference error.

  • Problem

    Planning objectives and the algorithmic tools that solve them are stated in different languages, leaving their shared decision variable and geometric relationship unresolved.

  • Method

    The paper embeds planning criteria in a resetting controlled Markov process and analyzes its stationary visitation measure using dual information geometry.

  • Results

    The visitation manifold has visitation-probability and log-policy affine charts that are dual under conditional entropy, making natural policy gradient and policy mirror descent one update and extending planning-as-inference to nonlinear free energies.

  • Takeaways & Limitations

    The visitation measure jointly supports variational inference, dynamic programming, and advantage estimation, while temporal-difference error acquires a marginal-utility interpretation.

  • Takeaways & Limitations

    A full measure-theoretic extension to continuous state and action spaces remains an open challenge, and the theoretical-neuroscience correspondence remains an open question.

Abstract

from arXiv · show

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.

1 Introduction

The paper addresses the mismatch between planning-as-inference objectives and dynamic-programming tools by treating a resetting process’s stationary visitation measure as the shared decision variable. It develops this variable’s dual geometry and uses it to extend planning beyond linear rewards.

  • Motivation: The visitation measure reconciles variational inference and dynamic programming as the stationary distribution over state–action events in a latent planning process.It is introduced as the object on which both planning-as-inference and dynamic programming operate.
  • Approach: A resetting planning process embeds the decision criterion into the dynamics, producing finite i.i.d. restart cycles and a normalized stationary visitation measure.The construction restarts from a fixed law at state-action-dependent rates.
  • Contributions: The visitation measure is the stationary law of the resetting process, and its Charnes–Cooper transform recovers the standard occupancy linear program across several reinforcement-learning criteria.The reset rate unifies discounted, finite-horizon, average-reward, and general early-termination criteria.
  • Contributions: Visitation probabilities and log-policies form Legendre-dual affine charts of a dually flat statistical manifold under conditional entropy.This gives a geometric account of natural policy gradient and policy mirror descent as the same update in different charts.
  • Contributions: Dual flatness extends planning-as-inference from linear rewards to nonlinear visitation functionals, with each iterate solved by one natural-gradient step.The framework also interprets temporal-difference error as an estimate of marginal utility.

2 The resetting planning processes

The paper embeds discounting and termination into a resetting controlled Markov process whose stationary state–action law is the visitation measure. This measure is simultaneously a geometric object, a variational posterior, and a unified reinforcement-learning occupancy representation.

  • Process construction: For a fixed memoryless policy, the resetting controlled chain has a stationary visitation measure νπ satisfying flow balance.Under the stated unichain condition, the stationary measure is unique and has full support.
  • Occupancy characterization: The Charnes–Cooper transformation converts the resetting fractional-linear program into the standard occupancy linear program for the corresponding reinforcement-learning criterion.The restart probability represents generalized discounting, with γ(s, a) = 1 − ρ(s, a).
  • Process construction: The resetting process ends an imagined life with probability ρ(s, a) and restarts from s0 ∼ µ, embedding the decision criterion into the dynamics.For constant ρ ≡ 1 − γ, the construction becomes the PageRank model of discounted Markov decision processes.
  • Visitation geometry: The visitation manifold is the relatively open feasible set of probability measures in the state–action simplex satisfying the flow constraint Kν = 0.Its dimension is |S|(|A| − 1), and it inherits Fisher–Rao/Hessian information geometry.
  • Inferential interpretation: The stationary visitation law is also a variational posterior over state–action events, so nonlinear free-energy objectives can be estimated from finite-length restart trajectories.Restart cycles are almost surely finite and independent across episodes.

3 Dually flat structure of the visitation manifold

The achievable visitation measures form a smooth statistical manifold with visitation probabilities and log-policies as dual affine coordinates. This dual flatness identifies natural-gradient and mirror-descent updates, while the centered advantage supplies their policy-space increment.

  • Manifold and charts: Visitation measures and policies are mutually invertible: conditioning a positive stationary measure gives the corresponding memoryless policy.The conditioning map preserves the induced policy under positive rescaling of the visitation measure.
  • Manifold and charts: The mixture coordinate is the visitation probability, while the exponential coordinate is the log-policy modulo state-wise constants.These coordinates are Legendre-dual under the conditional-negentropy potential.
  • Tangent structure: Tangent variations of visitation measures correspond to policy variations, with the flow constraint determining the marginal component from the conditional policy change.This identifies visitation and policy perturbations as two representations of the same tangent vector.
  • Cotangent structure: The centered advantage is the cotangent representative whose pairing with a visitation variation equals the directional derivative of return.It is invariant to potential-shaping and constant-gauge freedoms.
  • Metric structure: The Kakade metric is the visitation-weighted sum of per-state action-simplex Fisher metrics and is positive definite on the visitation tangent space.It is distinct from the joint Fisher–Rao metric of the visitation distribution, which does not respect the flow constraint.
  • Optimization: Mirror descent and natural policy gradient are the same update expressed in the mixture and exponential charts, respectively.For linear reward, the advantage is the exponential-coordinate increment; freezing it yields the explicit exponentiated-advantage approximation.
  • Scope: A full measure-theoretic extension to continuous state and action spaces remains an open challenge, while Wasserstein geometry provides a complementary structure.Wasserstein geometry transports visitation mass, unlike the Fisher–Rao/Hessian geometry’s reweighting in place.

4 Variational planning beyond linear rewards

The paper extends planning-as-inference from linear visitation rewards to nonlinear stationary free energies by using the dually flat visitation manifold. Each iteration becomes a linearized regularized MDP, while temporal-difference errors acquire a marginal-utility interpretation.

  • Generalization beyond linear rewards: Dually flat geometry extends planning-as-inference from linear visitation rewards to nonlinear stationary free energies.The framework covers maximum-state-entropy exploration, active inference, risk-sensitive objectives, and imitation learning.
  • Linear stationary free energy: The free-energy bound has slack equal to the Bregman divergence from the optimum, and it tightens as the visitation measure approaches that optimum.The log-partition function is the path-space log-evidence of the resetting generative model.
  • Linear stationary free energy: The stationary free-energy optimum satisfies a soft Bellman fixed point and reduces to textbook soft value iteration when ρ ≡ 1 − γ.The reset branch becomes a global constant in the advantage for constant resetting, recovering entropy-regularized MDP methods.
  • Generalization beyond linear rewards: Each nonlinear-planning iterate solves a linearized regularized MDP using the current functional gradient and an advantage update.Mirror descent supplies the per-iterate reward, and the implicit update solves the corresponding soft-Bellman fixed point.
  • Optimization and convergence: The policy sequence is a minorize–maximize envelope with a tabular and linear-MDP convergence guarantee under relative smoothness.The guarantee applies when the objective and geometric potential satisfy the required relative-smoothness condition; function approximation can break dual flatness and global convergence.
  • Theoretical neuroscience: The temporal-difference error is interpreted as an advantage for a shaped intrinsic reward whose magnitude slope is behavioral marginal utility.The neural interpretation is explicitly limited to the signal an optimal risk-sensitive agent would compute, not a claim about implementation.

5 Conclusion

The conclusion identifies visitation measures as a common object for variational inference, dynamic programming, and unbiased return estimation. It presents dual flatness as the basis for natural-gradient planning beyond linear rewards while leaving neural implementation and continuous-space extensions open.

  • Conclusion: Resetting visitation measures form a dually flat manifold with visitation probabilities and log-policies as dual affine charts.The charts are dual under conditional entropy.
  • Conclusion: Natural policy gradient and policy mirror descent are the same update expressed in the two affine charts.This links the geometric representation directly to optimization algorithms.
  • Conclusion: Nonlinear planning uses exact evidence lower bounds solved by one natural-gradient step per iterate.The conclusion states this as the extension of planning-as-inference beyond linear rewards.
  • Conclusion: Temporal-difference errors acquire an interpretation as marginal-utility estimates.The shared visitation manifold makes the variational, dynamic-programming, and return-estimation perspectives coincide.
  • Conclusion: Neural realization of the geometry and extension to continuous spaces remain future work.These are explicit scope boundaries in the conclusion.

A.1 Restarting visitations reproduce RL

The resetting formulation recovers standard reinforcement-learning criteria through a Charnes–Cooper transformation. The reset rate unifies discounted, finite-horizon, average-reward, state-action-dependent discounting, and early-termination objectives.

  • Restarting visitations reproduce RL: The Charnes–Cooper transformation converts the fractional-linear resetting formulation into the standard MDP linear program.The two programs have the same optimal value, with maximizers corresponding under mutually inverse maps.
  • Restarting visitations reproduce RL: A constant reset rate ρ ≡ 1 − γ recovers the discounted return.Under this choice, the normalization satisfies ⟨ρ, νπ⟩ = 1 − γ.
  • Restarting visitations reproduce RL: Clock-augmented resetting recovers the finite-horizon return by restarting when the horizon is reached.The reset condition is ρ(ŝ, a) = 1[t = H] with reset to t = 0.
  • Restarting visitations reproduce RL: As ρ ↘ 0, the visitation measure converges to the stationary law of the unreset chain when the policy-induced chain is unichain.The numerator gives the gain, while the ratio and transformation degenerate as the reset rate vanishes.
  • Restarting visitations reproduce RL: General reset rates represent state-action-dependent discounting and subsume early-termination objectives.The reset rate therefore carries the performance criterion inside the dynamics.

A.2 Equivalence of stationary and path-space free energies

The resetting-process construction makes the trajectory-level ELBO exactly equivalent to a stationary free energy, up to mean life-length. At the optimum, both formulations select the same policy, while state-action-dependent reset rates alter the Bellman backup.

  • Resetting process: The stationary visitation νπ is the normalized expected per-life occupancy of the resetting chain, with mean life-length Lπ = 1/⟨ρ, νπ⟩.The resetting process samples an initial state, follows the policy, and terminates according to state-action reset probabilities.
  • Free-energy equivalence: The path-space ELBO equals the stationary free energy multiplied by the mean life-length Lπ.The return and trajectory KL share the common factor Lπ, yielding log Zpath ≥ Lπ β−1F(νπ).
  • Optimal policy: The path-space and stepwise objectives have the same optimal policy because their relationship is a fixed positive scaling at the maximizer.The stepwise bound is the path-space bound per unit life-length.
  • Bellman interpretation: For constant ρ ≡ 1 − γ, the gain becomes the entropy-regularized start value and the gain-corrected equation reduces to the ordinary soft Bellman equation.With constant reset rate, c = β−1(1 − γ)⟨µ, V⋆⟩.

B.1 Concrete dual charts

The visitation manifold has explicit mixture and exponential charts: barycentric visitation coordinates describe a bounded polytope, while log-policy coordinates cover the full Euclidean space. These charts are Legendre dual under conditional negentropy.

  • Global chart geometry: Because the mixture chart is affine, its geodesics are straight lines and its mixture connection is flat.A convenient reference choice uses affinely independent deterministic-policy occupancies, making the polytope corners interpretable.
  • Mixture coordinates: The visitation manifold lies in the affine intersection of the fixed flow subspace ker K and the unit-mass hyperplane.Its tangent space is ker K ∩ ker 1⊤.
  • Mixture coordinates: The mixture chart uses barycentric weights over d+1 affinely independent visitation measures, giving a global coordinate system on the visitation polytope.The dimension is d = |S|(|A|−1), and the reduced weights provide the non-redundant coordinates.
  • Exponential coordinates: The exponential coordinate is the log-policy modulo state-wise constants, obtained as the gradient of the conditional-negentropy potential.The Legendre map θ = ∇φ(η) is invertible through η = ∇φ∗(θ).
  • Global chart geometry: The mixture image is a bounded polytope, whereas the exponential image is all of R^d.The Legendre map connects the bounded m-flat mixture family to the complete e-flat log-policy space.

B.2 The generalized Pythagorean theorem certifies dual flatness

A generalized three-point divergence identity expresses the Pythagorean deficit as the negative Fisher-metric pairing of mixture and exponential geodesics. The deficit vanishes exactly when those geodesics are orthogonal.

  • Orthogonality certificate: The exact Pythagorean relation provides a directly checkable certificate that the visitation manifold is dually flat.The canonical divergence is generated by the conditional-negentropy potential, with θ(ν) = log πν as exponential coordinate.
  • Geometric identity: The Bregman three-point identity decomposes the divergence difference into a pairing between the two geodesic chords meeting at the reference visitation.The mixture chord is ν1 − ν3, while the exponential chord is θ(ν2) − θ(ν3).
  • Orthogonality certificate: The Pythagorean deficit equals −gν3(ue, um) and vanishes exactly at metric orthogonality.The exponential chord is represented by the state-wise-centered log-policy difference, while the mixture chord lies in the mass-zero tangent space.

B.3 Linear MDPs and log-linear policies

In linear MDPs, visitation geometry reduces to a feature-expectation polytope paired with log-linear policies. Value-completeness makes the policy update finite-dimensional and exact, while continuous-space extensions remain conjectural.

  • Feature-expectation geometry: The reward- and dynamics-relevant visitation information is the d-dimensional feature-expectation polytope F = Φ⊤V+, where return is linear.Its deterministic-policy vertices recover the feature-matching geometry of apprenticeship learning.
  • Feature-expectation geometry: The Legendre-dual policy family is the log-linear class πw(a | s) ∝ exp⟨w, ϕ(s, a)⟩, with ϕ as sufficient statistic.The dual mixture coordinate is the feature expectation ψπ = Eνπ[ϕ].
  • Finite-dimensional updates: In linear MDPs, the exponentiated-advantage update collapses to a d-dimensional recursion because every advantage lies in span Φ.The resulting step is simultaneously policy mirror descent, Kakade’s natural policy gradient, and exact compatible function approximation.
  • Finite-dimensional updates: The update is exact for an exponential-family policy class precisely when its sufficient statistic is value-complete.When value-completeness fails, the update uses a projected advantage and incurs compatible-approximation error.
  • Scope boundary: The extension of dual flatness to genuinely continuous state and action spaces is stated as a conjecture under regularity conditions.The paper identifies compatibility between resetting flow constraints, measure-space Bregman geometry, and exponential-family dual flatness as the unresolved claim.

C Relative smoothness

The objective is relatively smooth with respect to the Fisher geometry on the visitation-measure manifold, with an explicit smoothness constant. Relative strong convexity fails because some policy variations change action probabilities without changing state transitions.

  • Relative smoothness: Proposition 1 establishes relative smoothness: −∇2F(ν) is bounded above by LF gν on every tangent space.The proof uses the second-order characterization of relative smoothness and the tangent-space constraints.
  • Proof mechanism: The entropy Hessian depends on the state-marginal variation as −∇2F(ν)(u, u) = τ∥uS∥2 1/νS, while the Fisher form uses the tangent decomposition into marginal and policy components.Flow balance determines the marginal component rather than leaving it free.
  • Proof mechanism: The flow identity gives uS = γ/(1−γ) R(P∗σ), linking marginal changes to conditional-policy variations through the state pushforward and resolvent operators.R is column-stochastic and maps the initial distribution to the discounted state occupancy.
  • Proof mechanism: χ2-contraction of the Markov operators controls the relevant weighted norms, while the reference-measure change contributes the concentrability factor C.The bound uses contractions for P∗ and R plus a separate weighted-norm comparison.
  • Relative smoothness: The resulting smoothness constant is LF = τγ2C/(1 −γ)2.The constant combines the entropy weight, discount factor, and concentrability term.
  • Failure of strong convexity: Relative strong convexity fails when two actions at a state share a transition kernel: an opposite action-probability variation can have zero objective curvature but positive Fisher curvature.For this variation, P∗σ = 0 and uS = 0, so −∇2F(ν)(u,u) = 0 while gν(u,u) > 0.
Loading 2609.04005v1…