The Dually Flat Geometry of Planning as Inference

The Dually Flat Geometry of Planning as Inference

arXiv:2609.04005v1 [cs.AI] 3 Sep 2026

Nikola Milosevic Neural Data Science and Statistical Computing Group Max Planck Institute for Human Cognitive and Brain Sciences Leipzig, 04103 Germany nmilosevic@cbs.mpg.de Asaki Kataoka Neural Computation Unit Okinawa Institute of Science and Technology Okinawa, 904-0495 Japan asaki.kataoka@oist.jp Nicolás Hinrichs Neural Data Science and Statistical Computing Group Max Planck Institute for Human Cognitive and Brain Sciences Leipzig, 04103 Germany hinrichsn@cbs.mpg.de Kenji Doya Neural Computation Unit Okinawa Institute of Science and Technology Okinawa, 904-0495 Japan doya@oist.jp Nico Scherf Neural Data Science and Statistical Computing Group Max Planck Institute for Human Cognitive and Brain Sciences Leipzig, 04103 Germany nscherf@cbs.mpg.de

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-asinference 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.

Preprint.

1

Introduction

Planning as inference reformulates the maximization of return as conditioning a generative model on the event of desirable outcomes. The reformulation places control on the same footing as perception and learning [25] and yields principled regularizers for control [41]. Active inference and the free energy principle extend it, connecting planning to the minimization of expected free energy [25, 17]. Yet while the free energy principle for perception is well developed, it is not settled how to carry it over to action selection [46], and a recurring difficulty is that the planning objective and the algorithmic tools that solve it efficiently are stated in different languages. The visitation measure of the agent’s latent planning process (the stationary probability of state–action events under a policy) reconciles the two: it is at once the variable of variational inference and the object on which dynamic programming acts. To make this precise we construct a resetting planning process, a latent controlled Markov chain that restarts from a fixed law at a state-action-dependent rate. Its i.i.d. restart cycles give unbiased estimates of reinforcement-learning returns, and its normalized stationary measure (the visitation measure) carries the geometry we study. Our contributions are the following. • A resetting characterization of the occupancy measure (§2): the visitation measure is the stationary law of the resetting process, and its Charnes–Cooper transform recovers the standard occupancy LP, unifying discounted, finite-horizon, average-reward, and general early-termination criteria through the reset rate. • The dually flat geometry of the visitation manifold (§3): the visitation measures form a dually flat statistical manifold [9] under the conditional entropy, with the visitation probabilities and the log-policies as Legendre-dual affine charts. This gives a single geometric account of KL control as covariant policy search [35, 14, 55, 40, 53, 45], in which the natural policy gradient and policy mirror descent are one update written in either chart [57]. • Variational planning beyond linear rewards (§4): dual flatness lets planning-as-inference generalize from linear rewards to nonlinear functionals of the visitation (convex-RL and active-inference free energies), where each iterate is an exact evidence lower bound solved by one natural-gradient step. The temporal-difference error becomes a natural-gradient estimate of the marginal utility of visitation, connecting dopaminergic prediction errors to stationary utility theory. Related work. Control and planning as inference dates back to Kalman’s observation that linearquadratic control and linear filtering are computationally equivalent [36], later termed Kalman duality [65, 37]. In robotics and machine learning, probabilistic reformulations of planning [66] and the maximum-entropy view of MDPs [71, 41] underlie much of contemporary deep reinforcement learning [2, 29, 28, 30]. The linear occupancy program approach to Markov decision processes has its roots in the operations research literature of the 1960’s [44, 20, 19, 24, 21]. A recent line of work beginning with maximumentropy exploration [31] replaces this linear objective by a nonlinear functional of the occupancy, unified by [70], and popularized by [69, 26, 50] as convex MDPs. A recurring subtlety, made precise by [51], is that these objectives generally differ between the finite-trial regime and the population limit. The resetting planning process is a normalized version of state-action-dependent discounting [67] which also has an LP formulation [33], and is precisely our transformed linear program [16]. The resetting is the PageRank random surfer, recently connected to policy optimization by [12]. The mechanism itself is discrete-time stochastic resetting [23]. The RL analogue of Amari’s natural gradient [7, 8] is the natural policy gradient (NPG), introduced by [35], justified information-geometrically by [14] and made an actor-critic method by [55]. Contemporary trust-region and proximal methods [58, 60, 2], are often interpreted as approximations, with connections to maximum-entropy methods [29], made precise by a convex optimization view [53, 27]. The dual mirror-descent viewpoint is policy mirror descent [40, 68]. We explain the geometric origin of this duality through a dual flatness perspective using the result of [57]. We build directly on the work of [49], who analyze policy gradient flows on normalized discounted occupancy measures, using the same Hessian geometry. 2

⟨ρ, y⟩=1

µ

νt-1

y

νt

νt+1

νt+2

···

ν

νπ

⟨1, ν⟩=1 ⊤ νt+1 = P⊤ ρ νt = (1−ρ) P νt + ρ µ

Kν π =0

(a) The induced visitation flow νt+1 = P⊤ ρ νt , a contraction with fixed point the stationary visitation ν π .

RS×A

0

(b) The flow cone {ν ≥ 0 : Kν = 0} and the two Charnes–Cooper slices y = ν/⟨ρ, ν⟩ and ν = y/⟨1, y⟩.

Figure 1: The restarting planning process and its visitation geometry. (a) Marginalising the reset chain for a fixed policy π gives the visitation flow νt+1 = Pρ⊤ νt , whose fixed point is the stationary measure ν π (Kν π =0), the object we treat as a point on a manifold and optimize in §3. (b) Relation of the standard RL liner program on occupancies y and the stationary linear fractional program on ν. Each ray is a set of scaling-equivalent solutions of Kν=0.

2

The resetting planning processes

We consider the model-based planning setting illustrated in Figure 1 a), where an agent simulates a Markov process to infer policies for action selection. The agent is equipped with a latent controlled Markov process which consists of a finite state space S, a finite action space A, and a resetting controlled Markov kernel Pρ (s′ |s, a) = (1 − ρ(s, a))P (s′ |s, a) + ρ(s, a)µ(s′ ),

(1)

where ρ : S × A → [0, 1] is a reset probability, µ ∈ P(S) is the start-state measure and P is the agent’s transition model of the environment. For constant ρ ≡ 1 − γ, this is the PageRank model of discounted Markov decision processes [12] with discount factor γ. For general ρ it is a normalized version of state-action-dependent discounting [67]. Fix a memoryless stochastic policy π : S → P+ (A) and define the state-action chain Pρπ (s′ , a′ | s, a) := π(a′ | s′ ) Pρ (s′ | s, a). If the reset process is unichain [56], the Pρπ admits a unique stationary visitation measure ν π ∈ P+ (S × A) of full support, satisfying the flow balance X X ν π (s, a) = Pρ (s | s′ , a′ ) ν π (s′ , a′ ), ∀s ∈ S. (2) a

s′ ,a′

In the following, we write the flow residual as the linear operator X X (Kν)(s) = ν(s, a) − Pρ (s | s′ , a′ ) ν(s′ , a′ )

(3)

s′ ,a′

a

and assume that ⟨ρ, ν⟩ > 0 for all ν ∈ V+ . The fractional-linear program ⟨r, ν⟩ , ν∈V+ ⟨ρ, ν⟩ sup

 V+ = ν > 0 : Kν = 0, ⟨1, ν⟩ = 1 ,

(4)

reproduces the standard performance criteria of reinforcement learning via its Charnes–Cooper transformation [16], which yields the standard linear program of the respective Markov decision process [44, 56]:  sup ⟨r, y⟩, Y+ := y > 0 : Ky = 0, ⟨ρ, y⟩ = 1 . (5) y∈Y+

The restart probability ρ then represents generalized discounting [67] and y is the standard occupancy measure of reinforcement learning for the general discount γ(s, a) = 1 − ρ(s, a), see Figure 1 b) and Appendix A.1. 3

im K⊤

a) ι

ker K

∂i ν

ν

∂i

P+ (S × A)

V+ b)

R |S|(|A|−1)

log π = ∇V+ φ(ν) visitation probabilities η(ν)

log-probabilities of the policy θ(ν)

mixture (m) chart

exponential (e) chart

Figure 2: The visitation geometry in one picture. a) the manifold via its inclusion ι(V+ ) = kerK ∩ P+ (S × A) in the state-action simplex. b) V+ admits two global affine charts, the mixture coordinate η of the visitation probabilities (left) and the exponential coordinate θ of the log-policies (right).

Motivation for the resetting planning process. Standard discounting attaches a factor γ t to a non-terminating trajectory. The resetting construction instead embeds the decision criterion into the process, which admits a natural interpretation of the occupancy measure in terms of infinite-horizon planning-as-inference as its stationary measure. With probability ρ(s, a) the agent’s imagined life time ends and restarts from s0 ∼ µ. The resulting normalization makes the same object interpretable in geometric and inferential terms: 1. Geometrically, a ν ∈ V+ is a probability measure on S × A, so it lives on a statistical manifold and inherits the Fisher–Rao / Hessian structure of information geometry. Precisely, the feasible set in Eq. 4 is the intersection of the open (S × A)-simplex with the affine subspace ker K, see Figure 2. Since the |S| rows of K have rank |S| − 1 and V+ is relatively open, it is a smooth submanifold of the open (S × A)-simplex of dimension |S|(|A| − 1). We will refer to V+ as the visitation manifold of the planning process. 2. Inferentially, the stationary law of the restarting process is a variational posterior where planning becomes conditioning a generative model on optimality, with ν the posterior over state–action events and the free energy its evidence bound (§4.1). With restarts, a single episode is an i.i.d. trajectory sample of a.s. finite length, and ν is the stationary measure of the process. Generalizing to nonlinear free energy functionals F (ν) is then meaningful from the point of view of variational inference and the descending directions can be estimated as policy advantages from a finite number of finite-length trajectories [3, 59, 50] without differentiation through the agent’s model. The visitation measure is thus the natural variable in which the information geometry, the control-asinference aspect, and the advantage estimator all can be jointly analyzed.

3

Dually flat structure of the visitation manifold

A point of the visitation manifold is a stationary visitation measure ν ∈ V+ . In the ambient linear coordinates, V+ is a relatively open convex subset of ker K, so these restrict to a global affine chart. The corresponding non-redundant mixture chart η(ν) ∈ R|S|(|A|−1) is constructed in Appendix B.1. We use ν and its ambient representation interchangeably for mixtures, and reserve η for statements needing a homeomorphism to R|S|(|A|−1) . The construction is visualized in Figure 2. 4

Visitations and policies are diffeomorphic. The maps (25) are mutually inverse real-analytic diffeomorphisms between the mass section V+ = {⟨1, ν⟩ = 1} and the reset-rate section {⟨ρ, y⟩ = 1} of the flow cone {z ≥ 0 : Kz = 0}, well defined since ⟨ρ, ν⟩ > 0. Every positive measure ν ∈ P+ (S × A) satisfying (2) is the P stationary measure of the memoryless policy recovered by conditioning, πν (a | s) = ν(s, a)/ a′ ν(s, a′ ), and conditioning is scale-invariant, πy = πν , so ν 7→ ν/⟨ρ, ν⟩ fixes the induced policy and moves only the normalization. The conditioning map y 7→ πy is a diffeomorphism [49], hence the same is true for ν 7→ πν . Dual affine coordinates.

The negative conditional entropy φ : V+ → R ∪ {+∞}, φ(ν) =

X

ν(s, a) , ′ a′ ν(s, a )

ν(s, a) log P

s,a

(6)

is convex on V+ (conditional entropy is concave in the joint law) and generates the statistical structure. Its differential is the exponential coordinate θ(ν) := ∇φ(ν),

∂φ = log πν (a | s), ∂ν(s, a)

(7)

the log-policy modulo state-wise constants cs . Thus the two affine charts of the same point ν are the mixture coordinate η (the visitation, m-flat) and the exponential coordinate θ (the log-policy, e-flat), Legendre-dual through θ = ∇φ(η) and η = ∇φ∗ (θ). Appendix B.1 derives this pairing on the explicit charts. Note that the reset rate deforms the m-chart and the Hessian geometry, but not the e-chart. This is the first form of the policy–visitation duality: a visitation and its log-policy are the mixture and exponential coordinates of one point. Tangent space. At ν ∈ V+ the tangent space is Tν V+ = ker K ∩ker 1⊤ , of dimension |S|(|A|−1). Every u ∈ Tν V+ decomposes uniquely into a marginal and a conditional part, P P u(s, a) = uS (s) πν (a | s) + νS (s) π̇u (a | s), (8) a π̇u (a | s) = 0, uS (s) := a u(s, a), and u 7→ π̇u is a linear isomorphism onto the tangent space of the product action simplex, a policy variation. Moreover, the marginal part is fully P determined by π̇u through the constraint Ku = 0, more precisely its linearized flow uS (s) = s′ ,a′ Pρ (s | s′ , a′ )u(s′ , a′ ). So a visitation variation u (m-side) and its policy variation π̇u (e-side) are two representations of one tangent vector, which is the infinitesimal form of the chart duality above. ∗ S×A ∗ Cotangent space. Dually, ) /(im K⊤ + R1), with P the ′cotangent ′space is Tν V+ = (R ⊤ (K λ)(s, a) = λ(s) − s′ Pρ (s | s, a)λ(s ). A reward function r and its shifts r + K⊤ λ + c1 represent the same cotangent vector, so covectors are state-action functions modulo (resetting-based) potential shaping (im K⊤ [54]) and constants (R1), exactly the symmetries of the advantage function [59]. They act as linear forms on tangent vectors on V+ .

The canonical representative of the cotangent class [r]ν is the centered advantage Aν [r], obtained by choosing λ = Vν [r], the value of the resetting chain. The value is defined by the resetting Bellman (Poisson) equation with its gain term, h i P ⟨r,ν⟩ Vν [r](s) + J(ν) ρ̄(s) = Eπν (·|s) r(s, a) + s′ Pρ (s′ | s, a) Vν [r](s′ ) , J(ν) = ⟨ρ,ν⟩ , (9) with ρ̄(s) = Eπν (·|s) [ρ(s, a)], which determines Vν [r] uniquely up to the constant gauge R1. The advantage is then the reduced cost   Aν [r] := r − J(ν) ρ − K⊤ Vν [r], Eπν (·|s) Aν [r](s, ·) = 0, (10) which is invariant to the constant gauge in Vν [r], so it is well defined as the centered representative of [r]ν even though Vν [r] is not.1 The advantage, read as the linear form ⟨Aν [r], ·⟩ on Tν V+ , is the cotangent object dual to visitation variations: for u ∈ Tν V+ the pairing ⟨Aν [r], u⟩ = ⟨[r]ν , u⟩ = dJν (u) is the directional derivative of the return, gauge-invariant because u ∈ ker K ∩ ker 1⊤ annihilates both im K⊤ and R1. 1 For constant ρ ≡ 1 − γ, ρ̄ ≡ 1 − γ and (10) reduces to the ordinary discounted advantage r + γP V [r] − E ν πν [·]: the

gain term J(ν)ρ = (1 − γ)J(ν)1 is a global constant and cancels in the centering.

5

The Hessian gν (u, v) = ∇2 φ(ν)(u, v) is X u(s, a) v(s, a) X uS (s) vS (s) X X π̇u (a | s) π̇v (a | s) gν (u, v) = − = νS (s) , ν(s, a) νS (s) πν (a | s) s,a s s a

Metric.

(11)

the ambient joint Fisher–Rao metric minus its state-marginal part, equivalently the visitation-weighted sum of per-state action-simplex Fisher-Rao metrics. This is the state-action Fisher metric of the natural policy gradient [35], also called the Kakade metric [49]. It is not the joint Fisher–Rao of ν, which does not respect the flow constraint. The Kakade metric is positive semidefinite on the ambient space (degenerate along policy-preserving marginal changes) but positive definite on Tν V+ , where those directions are excluded by Ku = 0. Dual flatness and mirror descent. The two charts each carry a flat connection. The mixture connection ∇m , for which ν is affine, has mixture geodesics t 7→ (1 − t)ν0 + tν1 . The exponential connection ∇e , for which θ = log πν is affine, has log-linear geodesics t 7→ π01−t π1t (normalized). They are mutually dual w.r.t. g, so (V+ , g, ∇m , ∇e ) is a dually flat statistical manifold [8, 63], with canonical divergence the visitation-weighted conditional KL, X
Dφ (ν∥ν ′ ) = νS (s) KL πν (· | s) ∥ πν ′ (· | s) , (12) s

the trust region of on-policy optimization. The m- and e-geodesics satisfy the generalized Pythagorean theorem (Appendix B.2). Consequently mirror descent and the natural policy gradient are one update in the two charts. With φ as mirror map, the proximal step νk+1 = arg min ⟨∇ν f (νk ), ν⟩ + α1 Dφ (ν∥νk ) ν∈V+

(13)

is, by [57, Thm. 1], natural-gradient descent in the exponential chart, and the two coincide for every α > 0. To first order in α they agree with the primal natural-gradient step νk − α g(νk )−1 ∇ν f (νk ), and all three coincide with the common Hessian gradient flow ν̇ = −gradg f as α → 0 [6]. Its firstorder condition is an additive step in the exponential coordinate. For f (ν) = −⟨r, ν⟩ the increment is the advantage, log πk+1 = log πk + α Aνk+1 [r] (mod values), (14) so the advantage is the exponential-coordinate increment. The update is implicit: the increment carries the advantage at the updated iterate νk+1 , through the value gauge that centers it. Freezing it to the current iterate, Aνk [r], gives the explicit exponentiated-advantage step πk+1 ∝ πk exp(αAνk [r]), the first-order approximation whose error is well understood [40, 68]. We proceed with the idealized implicit update. Infinite-dimesnsional and Wasserstein geometries. A full measure-theoretic extension to continuous S, A in the sense of [13] is an open challenge, but exponential families are known to form finite-dimensional dually flat submanifolds of the space of measures [9]. The analogue here is exponential-family policies in linearly parameterized MDPs [34], developed in Appendix B.3. V+ can also be equipped with an L2 -Wasserstein structure in continuous or discrete S × A [42] derived from the sample space’s own geometry. Unlike our Fisher–Rao/Hessian geometry, which reweights probability in place, Wasserstein geodesics transport visitation mass across states. The two are complementary natural gradients on the same polytope.

4

Variational planning beyond linear rewards

Planning as inference casts control as maximizing an evidence lower bound (ELBO) whose reward term is linear in the visitation ν ∈ V+ : variational inference on trajectory distributions with a structured posterior and exponential likelihood yields the standard soft-RL formulation, which we construct directly on the stationary measure of the restarting chain (Appendix A). Section 4.2 asks what happens when the objective is instead a nonlinear stationary free energy, inducing a generalutility MDP [70, 69]. This arises whenever a state-entropy term enters the free energy, which happens in maximum-state-entropy exploration [31, 10], active inference [48], risk-averse objectives [50], and imitation learning [32]. We show the dually flat geometry of V+ allows extending the soft-RL view to this case, reducing the nonlinear problem to a sequence of linear entropy-regularized MDPs. 6

at-2

at-1

at

s

st-1

st

st+1

st+2

o

ot-1

ot

ot+1

ot+2

y observed → perception

(a) Perception: infer q(x | y).

now t

at+1

µ

y expected → control

(b) Planning as inference on the resetting chain.

Figure 3: Planning as inference with restarting state visitation. (a) standard inferential perception infers latent state s from observation o; (b) the resetting generative model with restart measure µ. At each decision (st , at ), with prob. 1−ρ(st , at ) the state continues under P , with prob. ρ it restarts from µ. Planning is concerned with inferring the latent future at , e.g. as an open-loop sequence. 4.1

The linear stationary free energy

Fix a reference policy πref ∈ Π+ with stationary visitation ν ref ∈ V+ . A candidate π is the futureaction posterior and induces ν ∈ V+ (Figure 3), both stationary laws of the same resetting chain. Given a reward r and preference likelihood p(o=o⋆ | s, a) ∝ exp(β −1 r(s, a)), 0 < β < ∞, define the stationary free energy X
F(ν) := ⟨r, ν⟩ − β Dφ (ν ∥ ν ref ), Dφ (ν∥ν ref ) = νS (s) KL πν (· | s) ∥ πref (· | s) , (15) s

whose information cost is the manifold’s own canonical divergence (§3). Since ν ref has full support, the convex conjugate of Dφ (·∥ν ref ) at r/β,  log Z := sup ⟨r/β, ν⟩ − Dφ (ν∥ν ref ) , (16) ν∈V+

is finite, and Fenchel–Young gives the evidence lower bound, stated entirely in the visitation measure, β −1 F(ν) ≤ log Z,

log Z − β −1 F(ν) = Dφ (ν ∥ ν ⋆ ),

ν ⋆ = arg max F. V+

(17)

The bound tightens as ν → ν ⋆ , with slack the Bregman divergence to the optimum. log Z is a log-partition function, equal to the path-space log-evidence of the resetting generative model (Appendix A.2). The maximizer is a soft Bellman fixed point: stationarity of (16) gives the exponential-coordinate tilt log(π ⋆ /πref ) = β −1 Aν ⋆ [r] (mod values), with the value the fixed point of the log-sum-exp backup through the resetting kernel Pρ . For ρ ≡ 1 − γ the reset branch ρ⟨µ, V ⋆ ⟩ is a global constant that cancels in the advantage, so (17) reduces to textbook soft value iteration, π ⋆ ∝ πref exp(Aν ⋆ [r]/β): (15) is an entropy-regularized MDP [27] and standard soft-RL methods [41, 29, 2] apply with r as reward. The dually flat structure recasts this in geometric terms and, next, generalizes it to nonlinear free energies. 4.2

Nonlinear free energies and their solution

The linear free energy (15) is the special case where the information cost is the manifold’s own divergence. More generally, free-energy planning maximizes any smooth concave functional of the visitation, max −F (ν), F convex on V+ , (18) ν∈V+

the general-utility / convex-MDP regime [70, 69]. The objectives of interest are instances, differing only in ∇F : 7

(i) Biased generative models. Scoring states against a preferred distribution p̃ ∈ P+ (S) with base measure m(s, a) ∝ p̃(s)πref (a | s) gives −F (ν) = ⟨β −1 r, ν⟩ − KL(ν∥m), a biased generative model in the language of active inference [17]. Splitting KL(ν∥m) into its state and conditional parts, −F (ν) = ⟨rlin , ν⟩ − φ(ν) + H(νS ), rlin := β −1 r + log p̃ + log πref , (19) whose only nonlinear term is the state-marginal entropy H(νS ), the maximum-state-entropy exploration objective of [31]. (ii) Risk-sensitive objectives. A coherent or static risk functional of the sequential outcome, e.g. a CVaR-type measure, is likewise concave in ν [50, Table 1]. The nonlinearity appears in the preference rather than an exploration bonus, but the solution below is unchanged. Since −F is concave, (18) is a convex MDP and the soft Bellman picture of §4.1, which needed F linear in ν, no longer applies directly. Dual flatness resolves this by mirror descent, which needs only the linearization of F at the current iterate. Its ambient Euclidean gradient −∇F (νk ) fixes a covector representative of the differential −δF [νk ]/δν, and the mirror step νk+1 = arg max ⟨−∇F (νk ), ν⟩ − α1 Dφ (ν∥νk ) ν∈V+

is the per-iterate regularized MDP with that reward and advantage update   log πk+1 = log πk + α Aνk+1 − ∇F (νk ) (mod values).

(20)

(21)

The increment is the advantage of the frozen linearized reward, evaluated at the updated iterate. Convex free-energy planning is thus variational EM [11], the E-step the linear regularized MDP (20) ([27, MD-MPI]) and the M-step the relinearization ∇F (νk+1 ). The implicit update (21) is obtained by solving the soft-Bellman fixed point (27) with frozen reward −∇F , reference πk , and temperature α−1 . Its self-consistent value yields the next-iterate advantage Aνk+1 . Note that πk is the inner reference with temperature α and πref is the outer reference with temperature β. For (i), the intrinsic reward is the marginal utility of a stationary utility [39], −∇F (ν) = rlin − log ν(s, a) − 1 (using − log νS −log πν = − log ν), a policy-dependent novelty bonus on under-visited preferred states [69]. Assigning an ambient linear form to a differential on V+ is the Fisher-cometric construction of [52].2 The reset structure makes each per-iterate advantage estimable without backpropagation through the model, with tunable bias–variance [59, 3]. The policy sequence is the minorize–maximize envelope of F , as in deep actor-critics [58], except that function approximation breaks dual flatness and the global-convergence guarantee. In the tabular and linear-MDP cases (Appendix B.3) dual flatness supplies it whenever (F, φ) are relatively smooth: for −∇2 F ⪯ LF g on T V+ , mirror ascent with step 1/LF converges at rate LF /k in Dφ [43], and by §3 so does its NPG form. Relative smoothness must be checked per objective. We believe it holds broadly, since the generic obstruction is the state entropy, whose 1/νS curvature is matched by g ([48], finite horizon; Appendix C, discounted). 4.3

Application in theoretical neuroscience

The origin of reward prediction error. Decision making in the brain has long been understood through reinforcement learning and control as inference, most directly through the reward prediction error (RPE) encoded by phasic dopamine [62, 61]. A key refinement is that these errors are computed in units not of delivered reward but of a nonlinear utility of reward magnitude, recovered from risky-choice behaviour [64, 4] and later placed within distributional RL [18]. For sequential decisions, distributional utility measures and nonlinear occupancy programs can coincide: a coherent risk functional, e.g. a CVaR-type measure via its risk envelope, is a concave functional of the occupancy measure [50, Table 1], hence a convex MDP in the sense of §4. This gives a free-energy interpretation to the utility RPE through Koopmans’ stationary utility theory [39]. A stationary recursive utility takes the aggregator form F = V (r, δF ) relating immediate to prospective utility, but leaves V underdetermined. Our construction fixes it as the soft Bellman backup, the aggregator induced by the free-energy view. It is time-consistent because it is a stationary function of state. By this interpretation, the quantity dopamine mediates is the marginal utility of 2 Their target is the open simplex,

RS×A /(im K ⊤ + R1).

with cotangent quotient RS×A /R1.

8

Ours carries the flow gauge as well,

the stationary free energy, rk = −∇F (ν), whose reward-magnitude slope is the marginal utility u′ recovered behaviourally [64, 4]. The geometry predicts that the RPE is an estimate of the advantage that belongs to a shaped intrinsic reward, hence state- and policy-dependent rather than a fixed prediction target. This is a refinement that is testable in principle. Note however that this is an interpretive correspondence at the level of the signal an optimal risk-sensitive agent would compute, not a claim about neural implementation. Perception–action duality in the cortex. A parallel question is how control is realized in the cortex. Motivated by the shared laminar architecture of sensory and motor areas and by the classical filtering– control duality [36, 65], [22] maps inference and control variables onto (thalamo-)cortical circuitry, leaving the precise correspondence open. The information geometry suggests a sharper form of the question. A recent predictive-coding scheme derives cortical dynamics from variational free-energy minimization under an exponential-family assumption [38], a natural sensory-processing counterpart to our control-side framework. Both carry a dually flat manifold, ours pairing the stationary measure ν of the planning process with the log-policy under the conditional-negentropy potential φ, theirs pairing expectation and natural parameters under the log-partition. Whether cortex implements a sensorimotor duality then becomes the geometric question of what map relates these two dually flat manifolds, which we leave to future work.

5

Conclusion

In this early-stage work, we demonstrate that the stationary measures (visitation measures) of a resetting planning process form a dually flat statistical manifold, whose two affine charts are the visitation probabilities and the log-policies, dual under the conditional entropy. This singles out two optimization algorithms: the natural policy gradient and policy mirror descent are one update written in the two charts. Planning-as-inference then extends from linear rewards to nonlinear free energies with each iterate an exact evidence lower bound solved by one natural-gradient step, and the temporal-difference error acquires an interpretation as a marginal utility estimate. That the same manifold is simultaneously a variational posterior, a dynamic-programming object, and a source of unbiased per-cycle return estimates is what lets these views coincide. Whether neural circuits realize this geometry, and how it extends to continuous spaces, we leave to future work.

9

References [1] Pieter Abbeel and Andrew Y Ng. Apprenticeship learning via inverse reinforcement learning. In Proceedings of the twenty-first international conference on Machine learning, page 1, 2004. [2] Abbas Abdolmaleki, Jost Tobias Springenberg, Yuval Tassa, Remi Munos, Nicolas Heess, and Martin Riedmiller. Maximum a posteriori policy optimisation. arXiv preprint arXiv:1806.06920, 2018. [3] Alekh Agarwal, Sham M Kakade, Jason D Lee, and Gaurav Mahajan. On the theory of policy gradient methods: Optimality, approximation, and distribution shift. Journal of Machine Learning Research, 22(98):1–76, 2021. [4] Aydin Alikaya, Mackenzie Rack-Wildner, and William R Stauffer. Reward and value coding by dopamine neurons in non-human primates. Journal of Neural Transmission, 125(3):565–574, 2018. [5] Eitan Altman. Constrained Markov Decision Processes. Chapman & Hall/CRC, 1999. [6] Felipe Alvarez, Jérôme Bolte, and Olivier Brahic. Hessian riemannian gradient flows in convex programming. SIAM journal on control and optimization, 43(2):477–501, 2004. [7] Shun-ichi Amari. Natural gradient works efficiently in learning. Neural Computation, 10(2):251– 276, 1998. [8] Shun-ichi Amari. Information Geometry and Its Applications, volume 194 of Applied Mathematical Sciences. Springer, 2016. [9] Shun-ichi Amari and Hiroshi Nagaoka. Methods of information geometry, volume 191. American Mathematical Soc., 2000. [10] Yonatan Ashlag, Mirco Mutti, Aviv Tamar, and Kfir Yehuda Levy. Probing in the dark: State entropy maximization for POMDPs. In The Fourteenth International Conference on Learning Representations, 2026. [11] Pierre-Cyril Aubin-Frankowski, Anna Korba, and Flavien Léger. Mirror descent with relative smoothness in measure spaces, with application to sinkhorn and em. Advances in Neural Information Processing Systems, 35:17263–17275, 2022. [12] Konstantin Avrachenkov, Lorenzo Gregoris, and Nelly Litvak. Linking pagerank, time reversal, and policy evaluation. arXiv preprint arXiv:2605.00532, 2026. [13] Nihat Ay, Jürgen Jost, Hông Vân Lê, and Lorenz Schwachhöfer. Information geometry, volume 64. Springer, 2017. [14] J. Andrew Bagnell and Jeff Schneider. Covariant policy search. In Proceedings of the 18th International Joint Conference on Artificial Intelligence (IJCAI), pages 1019–1024, 2003. [15] Elliot Chane-Sane, Pierre-Alexandre Leziart, Thomas Flayols, Olivier Stasse, Philippe Souères, and Nicolas Mansard. Cat: Constraints as terminations for legged locomotion reinforcement learning. In 2024 IEEE/RSJ International Conference on Intelligent Robots and Systems (IROS), pages 13303–13310. IEEE, 2024. [16] Abraham Charnes and William W Cooper. Programming with linear fractional functionals. Naval Research logistics quarterly, 9(3-4):181–186, 1962. [17] Lancelot Da Costa, Thomas Parr, Noor Sajid, Sebastijan Veselic, Victorita Neacsu, and Karl Friston. Active inference on discrete state-spaces: A synthesis. Journal of Mathematical Psychology, 99:102447, 2020. [18] Will Dabney, Zeb Kurth-Nelson, Naoshige Uchida, Clara Kwon Starkweather, Demis Hassabis, Rémi Munos, and Matthew Botvinick. A distributional code for value in dopamine-based reinforcement learning. Nature, 577(7792):671–675, 2020. 10

[19] Guy De Ghellinck. Les problemes de decisions sequentielles. Cahiers du Centre d’Etudes de Recherche Opérationnelle, 2(2):161–179, 1960. [20] F d’Epenoux. Sur un probleme de production et de stockage dans l’aléatoire. Revue Française de Recherche Opérationelle, 14(3-16):4, 1960. [21] Cyrus Derman. Finite state Markovian decision processes. Academic Press, Inc., 1970. [22] Kenji Doya. Canonical cortical circuits and the duality of Bayesian inference and optimal control. Current Opinion in Behavioral Sciences, 41:160–167, October 2021. [23] Martin R Evans, Satya N Majumdar, and Grégory Schehr. Stochastic resetting and applications. Journal of Physics A: Mathematical and Theoretical, 53(19):193001, 2020. [24] Bennett Fox. Markov renewal programming by linear fractional programming. SIAM Journal on Applied Mathematics, 14(6):1418–1432, 1966. [25] Karl Friston. The free-energy principle: a unified brain theory? Nature reviews neuroscience, 11(2):127–138, 2010. [26] Matthieu Geist, Julien Pérolat, Mathieu Laurière, Romuald Elie, Sarah Perrin, Olivier Bachem, Rémi Munos, and Olivier Pietquin. Concave utility reinforcement learning: The mean-field game viewpoint. arXiv preprint arXiv:2106.03787, 2021. [27] Matthieu Geist, Bruno Scherrer, and Olivier Pietquin. A theory of regularized Markov decision processes. In Proceedings of the 36th International Conference on Machine Learning (ICML), volume 97 of Proceedings of Machine Learning Research, pages 2160–2169. PMLR, 2019. [28] David Ha and Jürgen Schmidhuber. Recurrent world models facilitate policy evolution. Advances in neural information processing systems, 31, 2018. [29] Tuomas Haarnoja, Aurick Zhou, Kristian Hartikainen, George Tucker, Sehoon Ha, Jie Tan, Vikash Kumar, Henry Zhu, Abhishek Gupta, Pieter Abbeel, et al. Soft actor-critic algorithms and applications. arXiv preprint arXiv:1812.05905, 2018. [30] Danijar Hafner, Timothy Lillicrap, Jimmy Ba, and Mohammad Norouzi. Dream to control: Learning behaviors by latent imagination. arXiv preprint arXiv:1912.01603, 2019. [31] Elad Hazan, Sham Kakade, Karan Singh, and Abby Van Soest. Provably efficient maximum entropy exploration. In Proceedings of the 36th International Conference on Machine Learning (ICML), volume 97 of Proceedings of Machine Learning Research, pages 2681–2691. PMLR, 2019. [32] Jonathan Ho and Stefano Ermon. Generative adversarial imitation learning. Advances in neural information processing systems, 29, 2016. [33] Héctor Jasso-Fuentes and Tomás Prieto-Rumeau. Constrained markov decision processes with non-constant discount factor. Journal of Optimization Theory and Applications, 202(2):897–931, 2024. [34] Chi Jin, Zhuoran Yang, Zhaoran Wang, and Michael I Jordan. Provably efficient reinforcement learning with linear function approximation. In Conference on learning theory, pages 2137– 2143. PMLR, 2020. [35] Sham M. Kakade. A natural policy gradient. In Advances in Neural Information Processing Systems 14 (NIPS), pages 1531–1538. MIT Press, 2001. [36] R. E. Kalman. A new approach to linear filtering and prediction problems. Journal of Basic Engineering, 82(1):35–45, 03 1960. [37] Hilbert J. Kappen, Vicenç Gómez, and Manfred Opper. Optimal control as a graphical model inference problem. Machine Learning, 87(2):159–182, 2012. [38] Asaki Kataoka and Kenji Doya. Extended predictive coding framework as variational freeenergy minimisation under exponential-family assumption, May 2026. 11

[39] Tjalling C Koopmans. Stationary ordinal utility and impatience. Econometrica: Journal of the Econometric Society, pages 287–309, 1960. [40] Guanghui Lan. Policy mirror descent for reinforcement learning: linear convergence, new sampling complexity, and generalized problem classes. Mathematical Programming, 198(1):1059– 1106, 2023. [41] Sergey Levine. Reinforcement learning and control as probabilistic inference: Tutorial and review. arXiv preprint arXiv:1805.00909, 2018. [42] Wuchen Li and Guido Montúfar. Natural gradient via optimal transport. Information Geometry, 1(2):181–214, 2018. [43] Haihao Lu, Robert M. Freund, and Yurii Nesterov. Relatively smooth convex optimization by first-order methods, and applications. SIAM Journal on Optimization, 28(1):333–354, 2018. [44] Alan S Manne. Linear programming and sequential decisions. Management Science, 6(3):259– 267, 1960. [45] Bianca Marin Moreno, Margaux Brégère, Pierre Gaillard, and Nadia Oudjane. Efficient modelbased concave utility reinforcement learning through greedy mirror descent. In International Conference on Artificial Intelligence and Statistics, pages 2206–2214. PMLR, 2024. [46] Beren Millidge, Alexander Tschantz, and Christopher L Buckley. Whence the expected free energy? Neural Computation, 33(2):447–482, 2021. [47] Nikola Milosevic, Leonard Franz, Daniel Haeufle, Georg Martius, Nico Scherf, and Pavel Kolev. Stochastic decision horizons for constrained reinforcement learning. arXiv preprint arXiv:2602.04599, 2026. [48] Nikola Milosevic, Nicolás Hinrichs, and Nico Scherf. Active inference as a convex markov decision process. arXiv preprint arXiv:2607.20152, 2026. [49] Johannes Müller and Guido Montúfar. Geometry and convergence of natural policy gradient methods. Information Geometry, 7(Suppl 1):485–523, 2024. [50] Mirco Mutti, Riccardo De Santi, Piersilvio De Bartolomeis, and Marcello Restelli. Challenging common assumptions in convex reinforcement learning. Advances in Neural Information Processing Systems, 35:4489–4502, 2022. [51] Mirco Mutti, Riccardo De Santi, and Marcello Restelli. The importance of non-markovianity in maximum state entropy exploration. In International Conference on Machine Learning, pages 16223–16239. PMLR, 2022. [52] Hiroshi Nagaoka. The fisher metric as a metric on the cotangent bundle. Information Geometry, 7(Suppl 1):651–677, 2024. [53] Gergely Neu, Anders Jonsson, and Vicenç Gómez. A unified view of entropy-regularized Markov decision processes, 2017. [54] Andrew Y Ng, Daishi Harada, and Stuart Russell. Policy invariance under reward transformations: Theory and application to reward shaping. In Icml, volume 99, pages 278–287. Citeseer, 1999. [55] Jan Peters and Stefan Schaal. Natural actor-critic. Neurocomputing, 71(7–9):1180–1190, 2008. [56] Martin L. Puterman. Markov Decision Processes: Discrete Stochastic Dynamic Programming. John Wiley & Sons, 1994. [57] Garvesh Raskutti and Sayan Mukherjee. The information geometry of mirror descent. IEEE transactions on Information Theory, 61(3):1451–1457, 2015. [58] John Schulman, Sergey Levine, Pieter Abbeel, Michael Jordan, and Philipp Moritz. Trust region policy optimization. In International conference on machine learning, pages 1889–1897. Pmlr, 2015. 12

[59] John Schulman, Philipp Moritz, Sergey Levine, Michael Jordan, and Pieter Abbeel. Highdimensional continuous control using generalized advantage estimation. arXiv preprint arXiv:1506.02438, 2015. [60] John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. arXiv preprint arXiv:1707.06347, 2017. [61] Wolfram Schultz. Dopamine reward prediction-error signalling: a two-component response. Nature reviews neuroscience, 17(3):183–195, 2016. [62] Wolfram Schultz, Peter Dayan, and P Read Montague. A neural substrate of prediction and reward. Science, 275(5306):1593–1599, 1997. [63] Hirohiko Shima. The geometry of Hessian structures. World Scientific, 2007. [64] William R Stauffer, Armin Lak, and Wolfram Schultz. Dopamine reward prediction error responses reflect marginal utility. Current biology, 24(21):2491–2500, 2014. [65] Emanuel Todorov. Efficient computation of optimal actions. Proceedings of the National Academy of Sciences, 106(28):11478–11483, 2009. [66] Marc Toussaint. Robot trajectory optimization using approximate inference. In Proceedings of the 26th International Conference on Machine Learning (ICML), pages 1049–1056, 2009. [67] Martha White. Unifying task specification in reinforcement learning. In International Conference on Machine Learning, pages 3742–3750. PMLR, 2017. [68] Lin Xiao. On the convergence rates of policy gradient methods. Journal of Machine Learning Research, 23(282):1–36, 2022. [69] Tom Zahavy, Brendan O’Donoghue, Guillaume Desjardins, and Satinder Singh. Reward is enough for convex MDPs. In Advances in Neural Information Processing Systems (NeurIPS), volume 34, 2021. [70] Junyu Zhang, Alec Koppel, Amrit Singh Bedi, Csaba Szepesvári, and Mengdi Wang. Variational policy gradient method for reinforcement learning with general utilities. In Advances in Neural Information Processing Systems (NeurIPS), volume 33, 2020. [71] Brian D. Ziebart. Modeling Purposeful Adaptive Behavior with the Principle of Maximum Causal Entropy. PhD thesis, Carnegie Mellon University, 2010.

13

A

Details on Planning as Inference

A.1

Restarting visitations reproduce RL

The fractional-linear functional

⟨r, ν π ⟩ (22) ⟨ρ, ν π ⟩ reproduces the standard performance criteria of reinforcement learning. To see this, we apply the Charnes–Cooper transformation of the fractional linear program J(π) =

⟨r, ν⟩ , ⟨ρ, ν⟩ ν∈V+

 V+ = ν > 0 : Kν = 0, ⟨1, ν⟩ = 1 ,

(23)

and obtain the standard linear program of Markov decision processes [44, 56]  sup ⟨r, y⟩, Y+ := y > 0 : Ky = 0, ⟨ρ, y⟩ = 1 ,

(24)

sup

y∈Y+

where ρ represents generalized discounting [67, 33]. The programs are equivalent in the sense that they share the same optimal value and their maximizers correspond under the mutually inverse maps y ν ν= (normalize), y= (unnormalize). (25) ⟨1, y⟩ ⟨ρ, ν⟩ The Charnes-Cooper transform in full would generically add another variable x with additional constraints {x : x ≥ 0, x = ⟨1, y⟩}. But because our main constraint Kν = 0 happens to be linear without affine part, the second variable does not enter the objective functional and its constraints do not interact with the others, so they are merely a definition and can be dropped. For ρ ≡  1 − γ one has ⟨ρ, ν π ⟩ = 1 − γ and (22) equals the discounted return Pconstant t E t≥0 γ r(st , at ) . The clock-augmented choice ρ(s̃, a) = 1[t = H] on states s̃ = (s, t) with reset to t = 0 recovers the finite-horizon return. The average-reward (gain) criterion arises in the vanishing-reset limit ρ ↘ 0: as the reset rate vanishes Pρπ → P π and ν π converges to the stationary law of the un-reset chain. The reward ⟨r, ν π ⟩ is the gain, provided P π is unichain so the limit is unique. Here the criterion is carried by the numerator alone, while the ratio J diverges and the transformation (25) degenerates as ⟨ρ, ν⟩ → 0. Finally, general ρ yields state-action-dependent discounting [67], and thus subsumes recent early-termination objectives from continuous control [15, 47]. A.2

Equivalence of stationary and path-space free energies

The linear free energy of §4.1 takes the form reward-minus-divergence as given. We certify that it is the evidence lower bound of a concrete generative model, and that Dφ is the standard trajectory KL. A life is one renewal cycle of the resetting chain: draw s0 ∼ µ. At each step draw at ∼ π(· | st ), flip a reset coin Rt ∼ Ber(ρ(st , at )), and continue to st+1 ∼ P (· | st , at ) if Rt = 0 or 1 end the life if Rt = 1. The length L is a.s. finite since E[L] = ⟨ρ,ν⟩ < ∞ and the process is unichain. Let p, q be the laws of one life under πref and π, with per-step preference likelihood −1 p(Ot =1 | st , at ) ∝ exp(β P r(st , at )). The expected per-life visitation is the (unnormalized) occupancy y π (s, a) = Eq [ t<L δ(st ,at ) (s, a)] = ν π (s, a)/⟨ρ, ν π ⟩ [56, 5], with mean life-length Lπ := ⟨1, y π ⟩ = 1/⟨ρ, ν π ⟩. P The two ELBOs. The path-space evidence is Zpath = Ep [exp(β −1 t<L r)], and Jensen against P q gives log Zpath ≥ Eq [β −1 t<L r] − KL(q∥p). Two identities collapse the right-hand side onto the stepwise free energy. First, the expected per-life return is the resetting objective, P  ⟨r,ν π ⟩ π Eq t<L r = ⟨r, y ⟩ = ⟨ρ,ν π ⟩ = J(π). Second, the trajectory KL is the canonical divergence: in log(q/p) the shared restart density µ(s0 ), P reset coins ρ, and transitions P cancel, leaving t<L log[π(at | st )/πref (at | st )]. Taking Eq and using y π (s, a) = ySπ (s)π(a | s), KL(q∥p) = KL(π∥πref ), y π = Lπ Dφ (ν π ∥ν ref ). 14

Both terms carry the common factor Lπ (the return is Lπ ⟨r, ν π ⟩∗ per life, the KL is Lπ Dφ ), so the path-space ELBO is exactly the stepwise free energy scaled by the mean life-length, log Zpath ≥ Lπ β −1 F(ν π ),

β −1 F(ν π ) = β −1 ⟨r, ν π ⟩ − Dφ (ν π ∥ν ref ).

(26)

π

At the maximizer L is a fixed positive scale, so log Zpath and the stepwise log Z of (16) share the same optimal policy; the stepwise bound (17) is (26) per unit life-length. The log-partition explicitly. The first-order conditions of (16) identify log Z with the multiplier c of the mass constraint ⟨1, ν⟩ = 1. Stationarity gives the tilt log(π ⋆ /πref ) = β −1 (r − K ⊤ V ⋆ ) − c1 = β −1 Aν ⋆ [r] − c1, where V ⋆ (the multiplier of Kν = 0, in reward units) is the soft bias and c the P soft gain; per-state normalization a π ⋆ (a | s) = 1 turns this into the gain-corrected soft Bellman equation   X  P V ⋆ (s)+βc = β log πref (a | s) exp β1 r(s, a)+(1−ρ(s, a)) s′ P (s′ | s, a)V ⋆ (s′ )+ρ(s, a)⟨µ, V ⋆ ⟩ , a

(27) solvable because the gain βc absorbs the constant null direction of the resetting kernel (cf. the gain-corrected Poisson equation of the cotangent space construction in the main text). It fixes V ⋆ up to the constant gauge, to which Aν ⋆ [r] = r − K ⊤ V ⋆ is invariant. Substituting the tilt back into (16) and using Kν ⋆ = 0, ⟨1, ν ⋆ ⟩ = 1, the two reward terms cancel and log Z = c = ⟨ρ, ν ⋆ ⟩ log Zpath =

log Zpath , Lπ ⋆

(28) ⋆

the soft gain: the per-life path log-partition log Zpath divided by the mean life-length Lπ = ⋆ 1/⟨ρ, ν ⋆ ⟩, i.e. the free-energy rate of the renewal process. For constant ρ ≡ 1 − γ, Lπ = 1/(1 − γ) and c = (1 − γ) log Zpath = β −1 (1 − γ)⟨µ, V ⋆ ⟩. For constant ρ ≡ 1 − γ the reset branch ρ⟨µ, V ⋆ ⟩ is a global constant, (27) reduces to the ordinary soft Bellman equation, and the gain is the entropyregularized start value c = β −1 (1 − γ)⟨µ, V ⋆ ⟩. For state-action-dependent ρ the reset term couples the backup through the scalar ⟨µ, V ⋆ ⟩.

B

Details on the Visitation Manifold

B.1

Concrete dual charts

First, we give the explicit mixture coordinate η. Because the reset kernel Pρ does not depend on the policy, the flow operator K is a fixed linear map, and the mixture chart is barycentric coordinates on a fixed affine space. The affine hull.

The point set V+ is a relatively open convex subset of the affine space n o H := ν ∈ RS×A : Kν = 0, ⟨1, ν⟩ = 1 ,

the intersection of the fixed flow subspace ker K with the sum-1-hyperplane. Since rank K = |S| − 1 (its rows sum to zero) and 1 is independent of those rows on ker K, H has dimension d := |S|(|A|−1), and its underlying linear space is the tangent space T = ker K ∩ ker 1⊤ of every point ν ∈ V+ . Note that ker 1⊤ on its own is the annihilator of constant functions and the standard tangent space of the open probability simplex [13, 52]. Mixture coordinates. Let d = |S|×(|A|−1) and fix d+1 affinely independent reference visitation measures ν0 , . . . , νd ∈ H. Every ν ∈ H has unique barycentric weights λi (ν) with ν=

d X

d X

λi (ν) νi ,

i=0

λi (ν) = 1,

i=0

and the mixture chart is the reduced weight vector
η(ν) := λ1 (ν), . . . , λd (ν) ∈ Rd , 15

(29)

with λ0 = 1 −

P

i≥1 λi and

η

−1

(x) = ν0 +

d X

xi (νi − ν0 )

(30)

i=1

its inverse. These maps give a chart for the mixture family of visitation measures and restricted to the point set, η is a homeomorphism of V+ onto the bounded open convex polytope η(V+ ) ⊂ Rd . A mixture geodesic is a measure segment t 7→ (1 − t)ν0′ + tν1′ . In the coordinate (29), it is the affine line t 7→ (1 − t)η(ν0′ ) + tη(ν1′ ), since η is affine. Hence its unique affine connection ∇m is flat in η and its Christoffel symbols vanish. η is a global ∇m -affine chart, as claimed in the main text. A convenient choice of the reference measures is d + 1 affinely independent deterministic-policy occupancies (vertices of the closure of V+ ), which makes the corners interpretable and η −1 a single matrix–vector product, the projection illustrated in Figure 2 on the bottom left. Exponential coordinates. The dual coordinate is the differential of φ read in the mixture chart. Since η 7→ ν(η) is affine (29) with constant Jacobian ∂ν/∂ηi = νi − ν0 , the chain rule gives, for the i-th dual coordinate, E D ∂φ ∂φ θi (ν) := , νi − ν0 = log πν , νi − ν0 , (31) = ∂ηi ∂ν where the last equality uses ∂φ/∂ν(s, a) = log πν (a | s) modulo state-wise constants: the ambient gradient of the conditional negentropy is the log-policy, and the state-wise constants annihilate the differences νi − ν0 (both are flows of unit mass, so (νi − ν0 ) has zero state marginals and pairs to zero against any state function cs ). Thus θ = ∇φ(η) is a linear image of the log-policy, and φ being strictly convex on V+ makes η 7→ θ a diffeomorphism onto its image; its inverse is η = ∇φ∗ (θ). In the intrinsic (basis-free) form this reads θ(ν) = log πν modulo state-wise constants—the exponential coordinate is the log-policy, and the reference-independent representative is fixed by choosing the gauge log πν (a◦s | s) = 0, giving θs,a = log πν (a | s) − log πν (a◦s | s). The two charts therefore have sharply different global shape. The mixture image η(V+ ) is a bounded polytope (an affine image of the bounded set V+ ), whereas the exponential image  θ(V+ ) = (log πν (a | s))a̸=a◦s = Rd is all of Rd : log-probabilities of an open simplex, modulo state-wise constants, range over the whole space. The Legendre map ∇φ is the homeomorphism between the bounded m-flat polytope and the complete e-flat space—the finite-dimensional shadow of a bounded mixture family paired with a complete exponential family. Its geodesics, the θ-affine lines, are the log-linear policy increments, and ∇e is flat in θ by the same argument that made ∇m flat in η. B.2

The generalized Pythagorean theorem certifies dual flatness

The dual flatness of V+ is equivalent to an exact Pythagorean relation for the Bregman divergence of φ, and this relation is a directly checkable certificate that the manifold carries the claimed structure. Throughout, Dφ is the canonical divergence of φ and θ(ν) = ∇φ(ν) = log πν the exponential coordinate. The starting point is an exact three-point identity that holds for any convex potential. Expanding each divergence through Dφ (ν∥ν ′ ) = φ(ν) − φ(ν ′ ) − ⟨θ(ν ′ ), ν − ν ′ ⟩ and cancelling the φ(ν1 ), φ(ν2 ) terms gives, for all ν1 , ν2 , ν3 ∈ V+ , Dφ (ν1 ∥ν2 ) − Dφ (ν1 ∥ν3 ) − Dφ (ν3 ∥ν2 ) = − θ(ν2 ) − θ(ν3 ), η(ν1 ) − η(ν3 ) .

(32)

The content of the correspondence is that the right-hand pairing is the metric inner product of the two geodesics meeting at ν3 , each read as the chord in its own affine chart. The mixture geodesic ν3 → ν1 is ν-affine, so its velocity is the ambient chord um = ν1 − ν3 ∈ Tν3 V+ ; the exponential geodesic ν3 → ν2 is θ-affine, so its velocity ue has constant θ-chord θ(ν2 ) − θ(ν3 ) = log πν2 − log πν3 . Since θ = ∇φ and g = ∇2 φ, the Legendre differential lowers ue to this chord, θ(ν2 ) − θ(ν3 ) = gν3 ue . The Euclidean pairing of this chord with um is then the metric pairing, because the marginal term of 16

g drops against a mass-zero tangent vector: log πν2 − log πν3 , um =

X ue (s, a) um (s, a) s,a

ν3 (s, a)

−

X ue,S (s) um,S (s) s

|

ν3,S (s) {z

= gν3 (ue , um ),

}

drops: log π is state-centered, um ∈ker 1⊤

(33) where the first equality uses that log πν2 − log πν3 is the state-wise-centered representative of ue /ν3 (log-conditionals carry no marginal component). Hence the Pythagorean deficit (32) equals −gν3 (ue , um ), vanishing exactly at g-orthogonality. B.3

Linear MDPs and log-linear policies

We generalize the dually flat structure to the linear MDP of Jin et al. [34], where dynamics and reward are linear in a known feature map ϕ : S × A → Rd , P (s′ | s, a) = ϕ(s, a), m(s′ ) ,

r(s, a) = ϕ(s, a), θr ,

with m = (m(1) , . . . , m(d) ) signed measures on S constrained so that ⟨ϕ(s, a), m⟩ ∈ P(S) for all (s, a) (i.e. Φm is row-stochastic), and θr ∈ Rd . Both sides collapse to d dimensions. Writing Φ ∈ RSA×d for the feature matrix and ψ π := Φ⊤ ν π = Eν π [ϕ] for the feature-expectation vector, the flow balance of §2 makes the state marginal affine in ψ π , νSπ (s′ ) = (1 − ρ) µ(s′ ) + γ ψ π , m(s′ ) . (34) The reward- and dynamics-relevant part of the occupancy is thus the d-dimensional linear image F := Φ⊤ V+ , a polytope on which the return is linear, J(π) = ⟨ψ π , θr ⟩, with deterministic-policy vertices. This is the feature-expectation polytope of apprenticeship learning/feature-matching [1]. Dually, the conjugate coordinate θ = log πν restricts to the log-linear class Πϕ = {πw (a | s) ∝ exp⟨w, ϕ(s, a)⟩}, and ϕ is forced as the sufficient statistic: for πw to be Legendre-dual to the mixture family with mean coordinate ψ = Eν [ϕ], the natural parameter w must pair with ψ through ϕ. The metric then pulls back to the d×d compatible-features Fisher matrix G(w) = Eν πw [Covπw ϕ] [35, 3]. The choice of conditional negentropy is what permits this, because it is degenerate along the statemarginal directions. Exact finite-dimensional PMD. In a linear MDP Qπ = ⟨ϕ, ω π ⟩ for every policy [34], so the advantage Aπ ∈ span Φ and the exponentiated-advantage step collapses to a d-dimensional recursion wk+1 = wk + α ω πk ,

G(wk )−1 ∇w J(πwk ) = ω πk ,

(35)

which is simultaneously the mirror step (13), Kakade’s natural policy gradient, and compatible function approximation [3] with zero approximation error (the compatible regression is exact because Aπk ∈ span Φ). More generally, an exponential-family class ΠT = {πw ∝ exp⟨w, T ⟩} is closed under the step iff T is value-complete, Aπ ∈ span T for all π ∈ ΠT . The linear MDP with T = ϕ is the canonical case. When value-completeness fails, (35) runs with ω πk the T -projection of Aπk and incurs the compatible-approximation error ∥Aπk − ΠT Aπk ∥, see [3]. Towards continuous spaces. The finite-dimensional collapse above suggests the dual flatness persists on genuinely continuous S, A for exponential-family policies, though a full treatment is beyond our scope. We record the expectation as a conjecture. Conjecture 1. Let S, A be standard Borel and ΠT = {πw ∝ exp⟨w, T ⟩} an exponential-family policy class with sufficient statistic T ∈ L2 and full-support restart law µ. Under the regularity conditions of parametrized measure models [13] (k-integrability of T , σ-finiteness of the occupancy) and with the conditional negentropy taken as Bregman potential in the directional-derivative sense ′ of measure-space mirror descent [11], the resetting visitation measures {ν πw : w ∈ Rd } form a finite-dimensional dually flat submanifold of P(S × A), with w 7→ log πw and w 7→ Eν πw [T ] the Legendre-dual coordinates and the compatible-features Fisher matrix as metric. In particular, the finite-dimensionality is inherited from T , not from |S||A|, and the mirror/NPG equivalence of §3 holds verbatim in the w-coordinates. 17

The two ingredients are independently established: finite-dimensional exponential families are dually flat submanifolds of the infinite-dimensional finite probability measure space [9, 13], and mirror descent with the entropy potential converges over measure spaces under relative smoothness [11]. So what the conjecture asserts is that the resetting flow constraint is compatible with both, i.e. that ν πw stays in the common domain where the directional-derivative Bregman geometry and the exponential-family dual flatness coincide.

C

Relative smoothness

Proposition 1 (Relative smoothness). F = ⟨r, ν⟩ + τ H(νS ) is LF -smooth relative to φ on V+ , i.e. −∇2 F (ν) ⪯ LF gν on Tν V+ for every ν ∈ relint V+ , with LF =

τ γ2 C , (1 − γ)2

C := sup ν∈V+

ν1 µ

∞

≤

1 , µmin

ν1 := P∗ ν.

(36)

Proof. By the second-order characterization of relative smoothness [43, Prop. 1.1], it suffices to show −∇2 F (ν)(u, u) ≤ LF gν (u, u) for every ν ∈ V+ and u ∈ Tν V+ = ker K ∩ ker 1⊤ . Two sides of the Hessian bound. The linear term of F drops and the Hessian of H(νS ) is −∇2 F (ν)(u, u) = τ ∥uS ∥21/νS the Fisher form on state-marginals. For the right side of the inequality, P P the tangent decomposition (8) u = uS πν +νS π̇ν substituted into gν (u, u) = s,a u2 /ν − s u2S /νS P cancels the marginal term and gives gν (u, u) = ∥σ∥21/ν with σ := νS π̇ν , a σ(s, a) = 0. It remains to bound ∥uS ∥21/νS by ∥σ∥21/ν . Marginal determined by the conditional (using u ∈ Tν V+ ). Differentiating the flow identity νS = (1 − γ)µ + γP∗ ν along u gives (I − γ(P π )⊤ )uS = γP∗ σ, whence γ uS = 1−γ R (P∗ σ),

R := (1 − γ)(I − γ(P π )⊤ )−1 ,

(37)

where R is column-stochastic with Rµ = νS (the discounted state occupancy), and P∗ the state pushforward. Without flow balance uS would be free but (37) ties it to σ. Both P∗ and R are Markov (row-/column-stochastic) operators, so each is a χ2 -contraction between the corresponding weighted norms (the data-processing inequality, cf. [13]): ∥P∗ σ∥21/ν1 ≤ ∥σ∥21/ν and ∥Rw∥21/νS ≤ ∥w∥21/µ . The single non-contraction is the reference change ∥w∥21/µ ≤ ∥ν1 /µ∥∞ ∥w∥21/ν1 , contributing the concentrability C (uniform since each ν1 is a probability vector, so C ≤ 1/µmin ). Combining,  2 2  γ γ ∥uS ∥21/νS = 1−γ ∥RP∗ σ∥21/νS ≤ 1−γ C ∥σ∥21/ν , which with the two sides above is the Hessian domination with LF = τ γ 2 C/(1 − γ)2 . Corollary 1 (Relative strong convexity fails). No η > 0 satisfies −∇2 F (ν) ⪰ η gν on Tν V+ : at any state s′ with two actions sharing a transition (P (·|s′ , a1 ) = P (·|s′ , a2 )), the variation π̇ν (a1 |s′ ) = +1, π̇ν (a2 |s′ ) = −1 has P∗ σ = 0, hence uS = 0 and −∇2 F (ν)(u, u) = 0 while gν (u, u) > 0.

18

Related documents

• Computational Geometry#3535
• Text-to-3D Policy: Fine-Grained Language-Behavior Alignment for Unseen Specification Generalization#1144566
• Home on the range: Planning and totality#745817
• Contact Modes Are Strata: What Geometric Structure Buys in Discrete-Continuous Planning#616700
• Contact Modes Are Strata: What Geometric Structure Buys in Discrete-Continuous Planning#616200