Conceptio › Archive › arXiv CS

arXiv CSopen access

Operator-Guided Invariance Learning for Continuous Reinforcement Learning

2026 · arxiv_cs

arXiv CS · Papers · License: Open Access · 2026

Open Source ↗Direct PDF ↓

knowledge-representation reasoning

artificial intelligence, reasoning, knowledge representation

Operator-Guided Invariance Learning for Continuous Reinforcement Learning

arXiv:2605.06500v1 [cs.LG] 7 May 2026

Zuyuan Zhang The George Washington University [email protected]

Fei Xu Yu The George Washington University [email protected]

Tian Lan The George Washington University [email protected]

Abstract Reinforcement learning (RL) with continuous time and state/action spaces is often data-intensive and brittle under nuisance variability and shift, motivating methods that exploit value-preserving structures to stabilize and improve learning. Most existing approaches focus on special cases, such as prescribed symmetries and exact equivariance, without addressing how to discover more general structures that require nonlinear operators to transform and map between continuous state/action systems with isomorphic value functions. We propose VPSD-RL (Value-Preserving Structure Discovery for Reinforcement Learning). It models continuous RL as a controlled diffusion with value-preserving mappings defined through Lie-group actions and associated pullback operators. We show that a value-preserving structure exists exactly when pulling back the value function and pushing forward actions commute with the controlled generator and reward functional. Further, approximate value-preserving structures with rigorous guarantees can be found when the Hamilton–Jacobi–Bellman mismatch is small. This framework discovers exact and approximate value-preserving structures by searching for the associated Lie group operators. VPSD-RL fits differentiable drift, diffusion, and reward models; learns infinitesimal generators via determining-equation residual minimization; exponentiates them with ODE flows to obtain finite transformations; and integrates them into continuous RL through transition augmentation and transformationconsistency regularization. We show that bounded generator/reward mismatch implies quantitative stability of the optimal value function along approximate orbits, with sensitivity governed by the effective horizon, and observe improved data efficiency and robustness on continuous-control benchmarks.

Introduction

Reinforcement learning (RL) has achieved great success across a variety of sequential decisionmaking tasks, yet modern RL are known to be sample-inefficient in complex environments and brittle under nuisance variability and shifts [Puterman, 2014, Sutton et al., 1998, Bertsekas, 2012, Kaelbling et al., 1996, Mnih et al., 2015, Schulman et al., 2015, 2017, Haarnoja et al., 2018, Xiang et al., 2019, Fujimoto et al., 2018, Kober et al., 2013, Levine et al., 2020, Zhang et al., 2026a,b,c, 2025a]. A key challenge is not optimization or exploration alone but also limited exploitation of underlying structures: if transformations exist to preserve (or nearly preserve) the optimal value function, valuelearning can generalize across otherwise distinct transitions/trajectories, significantly improving consistency, reducing complexity, and stabilizing learning [Dean and Givan, 1997, Ravindran and Barto, 2001, 2004, Li et al., 2006]. This viewpoint subsumes prior practice such as known forms Preprint.

of invariance-inducing structures (like symmetry, rotation, and linear scaling) and suggests a more ambitious goal: to discover general value-preserving structures and associated transformations–both exactly and approximately–from data and exploit them in reinforcement learning, while ensuring rigorous guarantees on the policy performance [Tobin et al., 2017, Laskin et al., 2020a, Yarats et al., 2021, Cohen and Welling, 2016, Bronstein et al., 2021]. This is particularly challenging in time- and state-continuous reinforcement learning, which is often modeled as controlled diffusion, a mathematically principled model widely used in robotics, finance, and physical processes [Oksendal, 2013, Karatzas and Shreve, 2014, Fleming and Soner, 2006, Yong and Zhou, 1999, Borkar and Borkar, 2008, Kushner, 1990]. In this model, optimality is governed by an operator—the controlled generator (or the induced semigroup)—and by the discounted Hamilton– Jacobi–Bellman (HJB) equation in viscosity form [Crandall et al., 1992, Bardi et al., 1997, Barles, 1994]. We define value-preserving transformations/mappings by Lie-group actions (due to the continuous spaces under consideration) and associated pullback operators. Consequently, we show that whether a transformation is “value-preserving” must be addressed through an operator-level question: Whether pulling back the value function (and pushing forward actions) commutes with the controlled generator and the reward functional, exactly or approximately. This operator-centric view provides a novel and tractable method to discover general value-preserving structures [Fleming and Soner, 2006, Bardi et al., 1997] in time- and space continuous RL. Existing approaches only considered certain known forms and special cases, such as exact symmetry by equivariant RL [Cohen and Welling, 2016, Kondor and Trivedi, 2018, Weiler and Cesa, 2019, Maron et al., 2018, Hutchinson et al., 2021, Bronstein et al., 2021] and invariant features by representation learning [Ha and Schmidhuber, 2018, Hafner et al., 2019, Laskin et al., 2020b, Stooke et al., 2021]. Data augmentation methods [Tobin et al., 2017, Laskin et al., 2020a, Yarats et al., 2021] can improve learning robustness if the relations are already known. Classical isomorphism and homomorphism formalize equivalence relations in MDPs, yet focusing only on discrete-time and discrete-state problems [Dean and Givan, 1997, Ravindran and Barto, 2001, 2004, Li et al., 2006]. Our proposed approach, VPSD-RL (Value-Preserving Structure Discovery for Reinforcement Learning), defines candidate transformations through an exact or approximate commutation property with respect to the controlled generator and reward function, using the discounted HJB operator [Fleming and Soner, 2006, Crandall et al., 1992, Bardi et al., 1997]. Algorithmically, VPSD-RL: (1) learns infinitesimal generators by minimizing determining-equation residuals; (2) exponentiates the learned vector fields via ODE flows to obtain finite transformations; and (3) injects these transformations into RL through transition augmentation and a transformation-consistency regularizer for value/policy networks [Olver, 1993, Bluman et al., 1989, Ibragimov, 2024, Lee, 2003]. A notable byproduct is that, when the discovered structure integrates to a (local) group action, it recovers classical Lie symmetries as a special case rather than a design premise [Olver, 1993, Bluman et al., 1989]. Theoretically, we prove that exact value-preserving structure implies exact value invariance of the ⋆ optimal value, while approximategenerator/reward mismatch yields the stability bound ∥V ◦

gϑ − V ⋆ ∥∞ ≤ β1 εr + εL ∥V ⋆ ∥C 2 . We further establish end-to-end guarantees for the VPSD-RL pipeline: stochastic structure discovery converges to approximate stationary points at standard SGD rates, learned infinitesimal generators are statistically consistent under realizability, numerical flow integration has order-hp error, and exact/approximate transition augmentation respectively preserves the Bellman fixed point or perturbs it by a controlled amount. VPSD-RL demonstrates improved data efficiency and robustness on representative continuous-control benchmarks [Fujimoto et al., 2018, Haarnoja et al., 2018].

Preliminaries

Discounted MDPs on Continuous State–Action Spaces. Let S and A be Borel subsets of Rd and Rm (or, more generally, Polish spaces equipped with their Borel σ-algebras). A discounted Markov decision process (MDP) is M = (S, A, P, r, γ) with γ ∈ (0, 1), where P (· | s, a) is a Markov transition kernel on S and r : S × A → R is a measurable reward bounded by |r(s, a)| ≤ Rmax . A (stationary Markov) kernel on A given s ∈ S. The value function hPpolicy π(· | s) is a stochastic i π π t of π is V (s) = E at ∼ π(· | st ), st+1 ∼ P (· | st , at ), and the t≥0 γ r(st , at ) s0 = s , optimal value is V ⋆ (s) = supπ V π (s). In finite-state MDPs, exact value-preserving correspondences 2

are often represented by permutations or homomorphisms. Because S, A are continuous, we work with smooth local transformation families on S and A. Lie-group actions provide one important special case and one convenient parameterization of transformation families used in our analysis, while stronger invertibility or closure properties can recover isomorphism-type correspondences within the broader value-preserving structure considered here. Controlled Diffusions, Semigroups, and Generators. Let S ⊆ Rd and A ⊆ Rm be open sets (or smooth manifolds with local coordinates). Consider infinite-horizon controlled diffusion in Itô form dst = b(st , at ) dt + Σ(st , at ) dWt ,

(1)

where Wt ∈ Rq is standard Brownian motion, b : S × A → Rd is the drift, and Σ : S × A → Rd×q is the diffusion coefficient. We use the same bounded reward r : S × A → R and fix a discount rate β > 0. For a (possibly stochastic) Markov control at ∼ π(· | st ), define the discounted return Z ∞ π π −βt V (s) = E e r(st , at ) dt s0 = s , V ⋆ (s) = sup V π (s). (2) π

For each fixed action a ∈ A, the associated infinitesimal generator La acts on f ∈ C 2 (S) by 1 Q(s, a) := Σ(s, a)Σ(s, a)⊤ . (La f )(s) = b(s, a) · ∇f (s) + tr Q(s, a) ∇2 f (s) , 2 Let (Pta )t≥0 be the Markov semigroup under constant control a: (Pta f )(s) = E f (st ) s0 = s, au ≡ a for u ∈ [0, t] .

(3)

(4)

Intuitively, Pta propagates functions forward in time under fixed dynamics, while La captures the d first-order infinitesimal change: formally, dt Pta f = La Pta f whenever the derivative is justified. Lie-Group Actions and Pullback Operators Let G be a Lie group acting smoothly on the state and action spaces through maps gθ : S → S and hθ : A → A for θ ∈ G, where each gθ is a diffeomorphism and each hθ is an invertible smooth map. We define the pullback operator Ugθ on functions f : S → R by (Ugθ f )(s) = f (gθ (s)). (5) We emphasize that we adopt the convention Ug f = f ◦ g (rather than f ◦ g −1 ). With this choice, (Ug f )(s) means: first move the state by the symmetry g, then evaluate f . This is the operator that naturally appears when expressing equivariance of generators/semigroups in pointwise form, e.g., comparing (Lh(a) (Ug f ))(s) with (Ug (La f ))(s). Viscosity HJB Basics and Well-Posedness Define the Hamiltonian operator on bounded continuous functions V : S → R by n o (HV )(s) = sup r(s, a) + (La V )(s) , (6) a∈A

and consider the discounted Hamilton–Jacobi–Bellman (HJB) equation βV (s) = (HV )(s).

(7)

Under the dynamic programming principle, the optimal value V ⋆ in (2) is characterized as the unique bounded viscosity solution of (7). We use viscosity solutions because V ⋆ need not be C 2 even when the coefficients are smooth; viscosity theory provides a robust notion of solution and a comparison principle that yields uniqueness. Theorem 2.1 (Existence and uniqueness of the optimal value). Under the HJB well-posedness conditions stated in Assumption B.9, the HJB equation (7) admits a unique bounded viscosity solution. Moreover, this solution coincides with the optimal value function V ⋆ defined in (2). This well-posedness is the formal anchor for our value-preservation analysis: once the transformed controlled generator and reward leave the HJB operator unchanged, uniqueness forces the optimal value function to inherit the corresponding value-preserving relation. Subsequent sections use this uniqueness-to-value-preservation logic repeatedly. 3

Value-Preserving Structure and Reduction for Controlled Diffusion

We characterize exact value-preserving structure of controlled diffusions at the level of the controlled generator (equivalently, the Markov semigroup), which directly links to the discounted HJB equation and avoids discretization dependence. Differentiating operator equivariance along one-parameter transformation families yields the determining equations, i.e., PDE constraints on drift, diffusion, and reward. This structure induces an intrinsic reduction via invariants, leading to value invariance along transformation orbits and structure-consistent optimal policies under mild conditions. 3.1

Exact Value-Preserving Structure

Definition 3.1 (Exact value-preserving structure). Let G be a Lie group acting smoothly on states and actions via (gϑ , hϑ ), ϑ ∈ G. We say (gϑ , hϑ ) is an exact value-preserving structure of the controlled diffusion if, for all ϑ ∈ G, all a ∈ A, and all f ∈ C 2 (S), La (f ◦ gϑ ) = (Lhϑ (a) f ) ◦ gϑ ,

La Ugϑ = Ugϑ Lhϑ (a) ,

equivalently

(8)

and r(gϑ (s), hϑ (a)) = r(s, a),

∀(s, a) ∈ S × A.

(9)

Equation (8) is the generator-level value-preservation condition: evolving after transforming (s, a) is equivalent, as seen through test functions, to transforming the outcome of evolution. In classical group-symmetry settings, this operator-commutation condition coincides with generator equivariance; here we use it as a criterion for value preservation. Crucially, since La depends on the diffusion only through Q = ΣΣ⊤ , the definition is intrinsic to the induced Markov process and does not depend on a particular factorization Σ. Lemma 3.2 (Semigroup form of exact value preservation). Assume for each a ∈ A that (Pta )t≥0 is a strongly continuous Markov semigroup on a Banach space F, and La is its generator with domain Dom(La ) ⊇ C 2 (S) ∩ F. Then Eq. 8 implies, for all t ≥ 0, h (a)

Pta Ugϑ = Ugϑ Pt ϑ

(10)

Conversely, if this semigroup identity holds for all t ≥ 0 and all f ∈ Dom(L on the corresponding generator domain.

hϑ (a)

), then Eq. 8 holds

The semigroup identity (10) compares the full time-t evolution under constant control. The generator form (8) is the infinitesimal version that is algebraically convenient for deriving determining equations and for interfacing with the HJB operator. 3.2

Determining Equations: From Global Value Preservation to Local PDE Constraints

To obtain explicit value-preservation constraints on (b, Q, r), we differentiate the operator identity in Definition 3.1 along a one-parameter subgroup. This yields linear PDE conditions, the determining equations, on the infinitesimal generators (X, Y ). Consider a one-parameter subgroup {ϑϵ }ϵ∈R ⊆ G with ϑ0 = e (the identity), and define the infinitesimal generators X(s) =

d gϑ (s) ∈ Ts S ∼ = Rd , dϵ ϵ ϵ=0

Y (a) =

d hϑ (a) ∈ Ta A ∼ = Rm . dϵ ϵ ϵ=0

(11)

Theorem 3.3 (Determining equations (necessity, Itô form under Ug f = f ◦ g)). Assume b, Q, r are differentiable and (gϑ , hϑ ) is an exact value-preserving structure in the sense of Definition 3.1. Then for all (s, a) ∈ S × A, X(s) · ∇s r(s, a) + Y (a) · ∇a r(s, a) = 0,

(12)

1 ∆Q(s,a) X(s) = 0, (13) 2 ∇s Q(s, a)[X(s)] − (∇s X(s)) Q(s, a) − Q(s, a) (∇s X(s))⊤ + ∇a Q(s, a)[Y (a)] = 0, (14) ∇s b(s, a) X(s) − (∇s X(s)) b(s, a) + ∇a b(s, a) Y (a) −

where ∇s Q(s, a)[X] denotes the directional derivative of the matrix field Q(·, a) at s along direction X(s), similarly for ∇a Q(s, a)[Y ], and ∆Q X := 4

⊤ Pd tr(Q ∇2s X1 ), . . . , tr(Q ∇2s Xd ) , tr(Q ∇2s Xk ) = i,j=1 Qij ∂ij Xk . In particular, if gϑϵ 2 is affine in s (equivalently ∇s X ≡ 0), then the Itô correction ∆Q X vanishes and (13) reduces to the first-order Lie-bracket form. Our value-preserving definition is posed at the level of the generator (8), which depends on the diffusion only through Q = ΣΣ⊤ . Therefore, the intrinsic necessary condition is (14). A condition written directly for Σ would be strictly stronger and generally not necessary, because different Σ factorizations can induce the same Q and hence the same generator. Theorem 3.4 (Local sufficiency). Assume b, Q, r are smooth. Let X and Y be smooth vector fields such that (12), (13), and (14) hold on an open set Ω ⊆ S × A. Assume the flows gϵ and hϵ generated d d by the ODEs dϵ gϵ (s) = X(gϵ (s)) and dϵ hϵ (a) = Y (hϵ (a)) exist and are diffeomorphisms for |ϵ| < ϵ0 . Then (gϵ , hϵ ) is an exact value-preserving structure on Ω for all |ϵ| < ϵ0 . Theorem 3.3 converts a global value-preserving operator identity into local PDE constraints on the coefficients. Theorem 3.4 states the converse: if the PDE constraints hold and the induced flows exist, then integrating the infinitesimal generators recovers local value preservation. Together, they justify treating value-preserving structure discovery as solving the determining equations. 3.3

Intrinsic Reduction by Invariants and Constraints

From value-preserving directions to reduced coordinates. Once value-preserving directions are characterized by vector fields X1 , . . . , XK , we can reduce the effective state dimension by passing to invariant coordinates that are constant along transformation orbits. This gives a coordinate-free description of the orbit space without explicitly constructing the quotient S/G. Let X1 , . . . , XK be linearly independent infinitesimal state generators and define the distribution D(s) = span{X1 (s), . . . , XK (s)} ⊆ Ts S.

(15)

A scalar function I : S → R is an invariant if it is constant along D: Xk (s) · ∇I(s) = 0,

k = 1, . . . , K.

(16)

Theorem 3.5 (Frobenius reduction and invariant coordinates). Assume D is a smooth distribution of constant rank K on a neighborhood, and it is involutive: [Xi , Xj ](s) ∈ D(s) for all i, j and all s in that neighborhood. Then locally there exist (d − K) functionally independent invariants I1 , . . . , Id−K whose joint map I(s) = (I1 (s), . . . , Id−K (s)) is constant on the integral leaves of D. In particular, the local orbit space can be represented intrinsically by the invariant coordinates z = I(s), without explicitly constructing a quotient S/G. Let the feasible set be Sfeas = {s : c(s) = 0, g(s) ≤ 0} for smooth constraints (c, g). A sufficient local condition for the value-preserving directions to preserve feasibility is tangency at feasible points: ∇c(s) Xk (s) = 0,

∇gℓ (s) Xk (s) ≤ 0 on active constraints (gℓ (s) = 0),

∀k, ℓ.

(17)

Under (17), the flow of each Xk stays in Sfeas for small time, so one can solve (16) restricted to Sfeas to obtain invariant coordinates for the post-constraint state space. 3.4

Value Invariance and Structure-Consistent Policies

From value preservation to control structure. We now connect the generator-level valuepreserving condition to consequences for optimal control. The key step is that exact value preservation leaves the HJB operator unchanged: reward compatibility and the generator-level condition imply H(Ugϑ V ) = Ugϑ (HV ), where invertibility of hϑ ensures the supremum over actions is preserved. Then, by uniqueness of bounded viscosity solutions (Theorem 2.1), the optimal value must be invariant along the transformation orbit. Theorem 3.6 (Exact value-preserving structure implies value invariance). Assume Assumption B.9 and let (gϑ , hϑ ) satisfy the exact structure conditions of Definition 3.1. Then the optimal value satisfies V ⋆ (gϑ (s)) = V ⋆ (s), ∀ϑ ∈ G, ∀s ∈ S. (18) Moreover, if I is a complete invariant coordinate map as in Theorem 3.5, then locally there exists V̄ ⋆ such that V ⋆ (s) = V̄ ⋆ (I(s)). 5

Value invariance already implies that the control problem effectively lives on the reduced coordinates z = I(s). To make this operational at the policy level, one often seeks a structure-consistent optimal selector. Corollary 3.7 (Orbit-consistent optimal actions (measurable selector form)). In addition to Assumption B.9, assume that for each s the maximizer set Argmax(s) := arg maxa∈A { r(s, a) + (La V ⋆ )(s) } is nonempty and admits a measurable selector a⋆ (s) ∈ Argmax(s). Then for any ϑ ∈ G, a ∈ Argmax(s) =⇒ hϑ (a) ∈ Argmax(gϑ (s)). In particular, one can choose an optimal selector satisfying the structure-consistency relation a⋆ (gϑ (s)) = hϑ (a⋆ (s)) on any neighborhood where such a selector exists. A concrete planar-rotation instance, together with its visualization, is provided in Appendix C.

Approximate Value-Preserving Structure and Stability of the Optimal Value

Exact value-preserving structure implies exact value invariance along transformation orbits (Section 3). Here we define approximate value-preserving structure by bounding the controlled-generator and reward mismatch over a test-function class, and show that it yields a corresponding approximate valueinvariance bound. The error scales linearly with the mismatch and is amplified by 1/β, reflecting increased sensitivity under weaker discounting. 4.1

Approximate Value-Preserving Structure as Bounded Operator Mismatch

Exact value-preserving equalities may be too strict in data-driven or partially observed systems. We therefore quantify value preservation through bounded mismatch of the transformed controlled generator and reward. Definition 4.1 ((εL , εr )-approximate value-preserving structure). Fix a test-function class F ⊆ C 2 (S) and the (standard) C 2 norm ∥f ∥C 2 := ∥f ∥∞ + ∥∇f ∥∞ + ∥∇2 f ∥∞ . We say (gϑ , hϑ ) is an (εL , εr )-approximate value-preserving structure if for all ϑ ∈ G, all a ∈ A, and all f ∈ F, La (f ◦ gϑ ) − (Lhϑ (a) f ) ◦ gϑ ∞ ≤ ϵL ∥f ∥C 2 .

(19)

r(gϑ (s), hϑ (a)) − r(s, a) ≤ εr .

(20)

and for all (s, a) ∈ S × A,

Condition (19) measures how far the transformed dynamics deviate from exact generator-level value preservation as seen through a function class F. The scaling by ∥f ∥C 2 is natural because La f involves first and second derivatives of f through b · ∇f and tr(Q∇2 f ). Condition (20) separately controls violations of reward compatibility. 4.2

From Approximate Value Preservation to Approximate Value Invariance

Under exact value preservation, reward compatibility and the generator-level value-preserving condition imply that the HJB operator H commutes with pullbacks: H(Ugϑ V ) = Ugϑ (HV ). With approximate value preservation, we obtain an inequality controlling the HJB commutator HUgϑ − Ugϑ H, which then translates into a bound on V ⋆ ◦ gϑ − V ⋆ via the discounted comparison principle. Theorem 4.2 (Approximate value-preserving structure implies approximate value invariance). Assume Assumption B.9. Let (gϑ , hϑ ) be an (εL , εr )-approximate value-preserving structure in the sense of Definition 4.1. Assume additionally that V ⋆ ∈ C 2 (S) ∩ F . Then for all ϑ ∈ G, 1 ∥V ⋆ ◦ gϑ − V ⋆ ∥∞ ≤ εr + εL ∥V ⋆ ∥C 2 . (21) β The bound (21) is intentionally stated in a strong-regularity regime (V ⋆ ∈ C 2 ) so that the generator mismatch in (19) can be evaluated on V ⋆ directly. In general, V ⋆ may only be a viscosity solution; extending (21) to that setting typically requires replacing ∥V ⋆ ∥C 2 by an appropriate smooth approximation argument (or a stability theorem for viscosity solutions under perturbations of the Hamiltonian). 6

Our VPSD-RL Algorithm

To develop VPSD-RL, we learn infinitesimal generators (X, Y ) by minimizing empirical determiningequation residuals from differentiable environment models, exponentiate them via ODE flows to obtain finite transforms (gα , hα ), and use these transforms for RL via (i) transition augmentation and (ii) transformation-consistency regularization. Algorithm 1 in Appendix gives the full procedure. Learning Infinitesimal Generators from Data Let b̂ω , Σ̂ω , r̂ω be differentiable parametric models of (b, Σ, r). Since the controlled generator depends on the diffusion only through Q = ΣΣ⊤ , we work with the induced diffusion matrix model Q̂ω (s, a) := Σ̂ω (s, a)Σ̂ω (s, a)⊤ ∈ Rd×d . We parameterize infinitesimal generators by Xθ : S → Rd and Yϕ : A → Rm . Motivated by Theorem 3.3, define residuals Rr (s, a) := Xθ (s) · ∇s r̂ω (s, a) + Yϕ (a) · ∇a r̂ω (s, a),

(22) 1 Rb (s, a) := ∇s b̂ω (s, a)Xθ (s) − (∇s Xθ (s))b̂ω (s, a) + ∇a b̂ω (s, a)Yϕ (a) − ∆Q̂ω (s,a) Xθ (s). 2 (23)

RQ (s, a) := ∇s Q̂ω (s, a)[Xθ (s)]

(24) ⊤

− (∇s Xθ (s)) Q̂ω (s, a) − Q̂ω (s, a)(∇s Xθ (s)) + ∇a Q̂ω (s, a)[Yϕ (a)].

(25)

Here ∇s Q̂ω (s, a)[X] denotes the directional derivative of the matrix field Q̂ω (·, a) at s along direction Xθ (s), and similarly for ∇a Q̂ω (s, a)[Y ]. In the experiments, the generator class is restricted to affine or locally linear vector fields, so the Itô correction term ∆Q̂ω Xθ vanishes. For general nonlinear generator classes, Rb should include the additional − 12 ∆Q̂ω Xθ term from Theorem 3.3. We minimize a weighted squared residual loss over a replay distribution ρ(s, a): h i 2 min E(s,a)∼ρ ∥Rb (s, a)∥22 + λQ ∥RQ (s, a)∥2F + λr |Rr (s, a)|2 + λnrm Es∼ρS ∥Xθ (s)∥22 − 1 , θ,ϕ

(26) where ρS is the state marginal of ρ. The normalization penalty prevents the trivial zero-field solution and fixes the overall scale of Xθ . (When learning multiple generators, one may additionally impose orthonormality constraints E[Xi · Xj ] = δij to fix a basis of the symmetry algebra.) The residuals in (22) enforce the infinitesimal determining equations. After exponentiation (Section 5), small residuals imply that the finite-step transforms (gα , hα ) yield small generator mismatch for small |α| (up to model and numerical integration errors), thereby providing a concrete path from minimizing (26) to reducing (εL , εr ) in Definition 4.1. Exponentiation via Flow Integration Given learned infinitesimal generators (Xθ , Yϕ ), we obtain finite transformations by integrating their flows (i.e., exponentiating in the Lie sense): d d gα (s) = Xθ (gα (s)), g0 (s) = s, hα (a) = Yϕ (hα (a)), h0 (a) = a. (27) dα dα In practice, we integrate (27) numerically for small |α|. If S or A is constrained (e.g., box constraints), we apply a projection/retraction after each solver step to keep the iterates feasible. Value-Preserving Reinforcement Learning Updates. Once a finite transformation (gα , hα ) is available, we use it in RL in two complementary ways. First, transition augmentation maps collected samples (s, a, s′ , r) to transformed samples (gα (s), hα (a), gα (s′ ), r), allowing the replay buffer to share experience along discovered valuepreserving orbits. Second, transformation-consistency regularization encourages the learned value and policy networks to make compatible predictions on transformed state-action pairs. Under exact value preservation these operations are structure-preserving data-sharing steps; under approximate value preservation their effect is controlled by the mismatch bounds in Section 4. Detailed losses and implementation choices are deferred to Appendix A.1. Algorithm 1 in Appendix alternates among (i) data collection with the current policy, (ii) optional environment model fitting to obtain differentiable b̂ω , Q̂ω , r̂ω , (iii) value-preserving structure discovery by minimizing (26) using automatic differentiation, and (iv) RL updates using augmentation and/or transformation-consistency regularization with the induced finite transforms. 7

Determining-residual convergence

0.10 0.05 0.00 0

60K

120K

180K

Env steps

240K

(a) Rot2D

300K

PostConstraint-Rot2D-v0 Ref2D-DoubleWell-v0 Rot2D-Point-v0

0.3 0.2

SymNav-v1

PostConstraint-Rot2D-v0 Ref2D-DoubleWell-v0 Rot2D-Point-v0

0.4

250 500

Eval return

0.15

RQ2

Rr2

0.20

0.35 0.30 0.25 0.20 0.15 0.10 0.05 0.00

Rb2

0.25

Value non-invariance |V(s) V(g (s))|

Alignment between learned and ground-truth generators

PostConstraint-Rot2D-v0 Ref2D-DoubleWell-v0 Rot2D-Point-v0

0.30

750 VPSD-RL SAC TD3 PPO DDPG A2C DynRand RARL EPOPT

1000 1250 1500

0.1 0

60K

120K

180K

Env steps

240K

1750

300K

(b) Q-res.

60K

120K

180K

Env steps

240K

300K

2000

60K

120K

180K

Env steps

240K

300K

(d) SymNav-V1

Figure 1: Diagnostics on interpretable controlled diffusions and one representative SymNav-15 variant. The first panels show determining-residual convergence, generator alignment, and value non-invariance diagnostics for systems with known value-preserving transformations. The last panel shows a representative learning curve on SymNav-V1. Curves report mean ± standard error over seeds.

Guarantees of VPSD-RL

We establish an end-to-end guarantees for VPSD-RL which connects (i) learning decision-preserving infinitesimal generators from data, (ii) exponentiating them into finite transformations, and (iii) injecting these transformations into RL. The full formal statements (assumptions, theorems, and proof sketches) are provided in Appendix B. Optimization guarantee. Theorem B.2 shows that stochastic optimization of the structure-discovery objective reaches approximate stationary points under standard smoothness and bounded-variance PT −1 2 Lsym (θ0 ,ϕ0 )−Linf 1 2 2 conditions: T t=0 E ∥∇Lsym (θt , ϕt )∥2 ≤ + ηLσ , and diminishing stepηT √ sizes give the standard O(1/ T ) stationarity rate. Statistical consistency. Theorem B.4 states that, under realizability and consistent environmentmodel estimation, learned generators converge in ρ-mean to solutions of the determining equations. Consequently, the induced generator and reward mismatches satisfy εL → 0 and εr → 0 on the data support. Exact augmentation. Theorem B.7 gives a discrete-time contraction analogue: if the augmentation is exactly value-preserving, Bellman backups computed from augmented samples preserve the same optimal Bellman operator, so tabular value iteration keeps the usual convergence rate ∥Vk − V ⋆ ∥∞ ≤ γ k ∥V0 − V ⋆ ∥∞ . Approximate augmentation. Theorem B.8 shows that if augmentation is only (εP , εr )-approximate, the perturbed Bellman operator remains a γ-contraction and its fixed point satisfies ∥Ṽ ⋆ − V ⋆ ∥∞ ≤ 1 ⋆ 1−γ (εr + γ εP ∥V ∥∞ ) . The numerical flow approximation guarantee, which bounds ∥ĝα − gα ∥ and ∥ĥα − hα ∥ by order hp , is stated in Theorem B.6 in the appendix.

Experiments

Our experiments provide supporting evidence for the value-preserving structure theory in Sections 3–6. We evaluate whether (i) infinitesimal generators can be recovered from data through determiningequation residual minimization, (ii) the induced finite transformations yield small operator mismatch and value non-invariance, and (iii) using the discovered transformations improves sample efficiency and robustness in continuous-control RL. We report episodic return as a function of environment steps, final performance over the last evaluation window, area under the learning curve (AUC) up to a fixed budget, and wall-clock cost. For interpretable systems with known ground-truth generators, we additionally report residual statistics E∥Rb ∥22 , E∥RQ ∥2F , E|Rr |2 , generator alignment, and empirical value non-invariance Es,α |V (s) − V (gα (s))|. Full environment descriptions, hyperparameters, ablations, and all extended curves are deferred to Appendix D. Figure 1 supports the mechanism predicted by the theory. On interpretable controlled diffusions, minimizing the determining-equation residuals recovers transformation directions aligned with the known generators and reduces empirical value non-invariance along the learned flow. On the representative SymNav-15 task, using the learned transformations improves sample efficiency relative 8

MJ-Hopper-v4

2000 1500

4000

1000

3000

6000 4000

Eval return

2500

MJ-Ant-v4

VPSD-RL SAC TD3 PPO DDPG A2C DynRand RARL EPOPT

5000

Eval return

3000

Eval return

MJ-Walker2d-v4

VPSD-RL SAC TD3 PPO DDPG A2C DynRand RARL EPOPT

3500

2000

1000

500

2000

0 0

1.0M

2.0M

3.0M

Env steps

4.0M

5.0M

(a) Hopper-v4

VPSD-RL SAC TD3 PPO DDPG A2C DynRand RARL EPOPT

2000

1.0M

2.0M

3.0M

4.0M

Env steps

5.0M

(b) Walker2d-v4

1.0M

2.0M

3.0M

Env steps

4.0M

5.0M

Figure 2: MuJoCo locomotion performance on Hopper-v4, Walker2d-v4, and Ant-v4. Curves report mean ± standard error over seeds.

Table 1: Compact quantitative summary on Ant-v4. Numbers are mean ± std over seeds. Final return is averaged over the last evaluation window; AUC is computed over a fixed step budget; wall-clock is seconds per 105 environment steps.

Method

Final return ↑

AUC@B ↑

Wall-clock ↓

A2C DDPG PPO SAC TD3

-790.30±103.76 2637.74±1385.82 1922.56±615.62 4477.82±677.98 3984.84±1111.78

-388.66±127.69 924.95±307.18 1414.10±205.77 1227.46±172.73 1463.37±123.77

391±113 437±8 215±0 397±159 767±77

DynRand RARL EPOpt

5149.99±110.37 4903.49±615.20 3872.36±89.81

1440.34±152.56 1558.29±227.95 1533.66±254.87

434±108 619±119 343±81

VPSD-RL (Aug+Reg)

5822.15±470.25

1653.21±336.34

471±161

to standard continuous-control baselines, suggesting that the recovered structure is useful beyond the synthetic diagnostic setting. Figure 2 shows that the benefit of VPSD-RL is task-dependent but consistent with the amount of reusable structure in the task. The gain is modest on Hopper-v4, whose dominant structure is a single hopping template; clearer gains appear on Walker2d-v4, where left-right phase correspondence creates reusable coordination; and the strongest improvement appears on Ant-v4, where multi-limb coordination produces richer approximately value-preserving transformation patterns. Thus the structural summary previously reported as a separate table can be stated directly in text. The compact comparison in Table 1 shows that VPSD-RL improves final return and AUC on Ant-v4 while adding moderate computational overhead. Detailed aggregate results for all SymNav-15 variants, full MuJoCo curves, robustness baselines, and ablations are provided in Appendix D. These empirical diagnostics are used as reliability checks for the discovered transformations, rather than as an assumption that every continuous-control task admits a useful global symmetry.

Conclusion

We introduced VPSD-RL, an operator-level framework for discovering and exploiting valuepreserving structure in continuous-time reinforcement learning. Rather than assuming a prescribed symmetry, VPSD-RL defines exact and approximate value preservation through commutation of the transformed controlled generator and reward with the HJB operator, making classical Lie symmetries a special case. We showed that exact value-preserving structure implies value invariance, while bounded generator/reward mismatch yields quantitative stability of the optimal value along approximate orbits. We then turned these principles into a practical pipeline that learns infinitesimal generators from data, exponentiates them into finite transformations, and uses them for transition augmentation and transformation-consistency regularization. The resulting theory and experiments suggest that value-preserving structure discovery can improve data efficiency and robustness in continuous-control RL, especially when useful transformations are present but not known in advance. These gains should be interpreted under the stated local-modeling assumptions: the learned flows are local and their reliability is assessed through residual and invariance diagnostics. Limitations and broader impacts are discussed in Appendix. 9

References Martin L Puterman. Markov decision processes: discrete stochastic dynamic programming. John Wiley & Sons, 2014. Richard S Sutton, Andrew G Barto, et al. Reinforcement learning: An introduction, volume 1. MIT press Cambridge, 1998. Dimitri Bertsekas. Dynamic programming and optimal control: Volume I, volume 4. Athena scientific, 2012. Leslie Pack Kaelbling, Michael L Littman, and Andrew W Moore. Reinforcement learning: A survey. Journal of artificial intelligence research, 4:237–285, 1996. Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Andrei A Rusu, Joel Veness, Marc G Bellemare, Alex Graves, Martin Riedmiller, Andreas K Fidjeland, Georg Ostrovski, et al. Human-level control through deep reinforcement learning. nature, 518(7540):529–533, 2015. 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. John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. arXiv preprint arXiv:1707.06347, 2017. Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with a stochastic actor. In International conference on machine learning, pages 1861–1870. Pmlr, 2018. Jiaqi Xiang, Qingdong Li, Xiwang Dong, and Zhang Ren. Continuous control with deep reinforcement learning for mobile robot navigation. In 2019 Chinese Automation Congress (CAC), pages 1501–1506. IEEE, 2019. Scott Fujimoto, Herke Hoof, and David Meger. Addressing function approximation error in actor-critic methods. In International conference on machine learning, pages 1587–1596. PMLR, 2018. Jens Kober, J Andrew Bagnell, and Jan Peters. Reinforcement learning in robotics: A survey. The International Journal of Robotics Research, 32(11):1238–1274, 2013. Sergey Levine, Aviral Kumar, George Tucker, and Justin Fu. Offline reinforcement learning: Tutorial, review, and perspectives on open problems. arXiv preprint arXiv:2005.01643, 2020. Zuyuan Zhang, Sizhe Tang, and Tian Lan. Cochain perspectives on temporal-difference signals for learning beyond markov dynamics. arXiv preprint arXiv:2602.06939, 2026a. Zuyuan Zhang, Mahdi Imani, and Tian Lan. Geometry of drifting mdps with path-integral stability certificates. arXiv preprint arXiv:2601.21991, 2026b. Zuyuan Zhang, Zeyu Fang, and Tian Lan. Structuring value representations via geometric coherence in markov decision processes. arXiv preprint arXiv:2602.02978, 2026c. Zuyuan Zhang, Hanhan Zhou, Mahdi Imani, Taeyoung Lee, and Tian Lan. Learning to collaborate with unknown agents in the absence of reward. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 39, pages 14502–14511, 2025a. Thomas Dean and Robert Givan. Model minimization in markov decision processes. In AAAI/IAAI, pages 106–111, 1997. Balaraman Ravindran and Andrew G Barto. Symmetries and model minimization in markov decision processes, 2001. Balaraman Ravindran and Andrew G Barto. Approximate homomorphisms: A framework for non-exact minimization in markov decision processes. 2004. Lihong Li, Thomas J Walsh, and Michael L Littman. Towards a unified theory of state abstraction for mdps. AI&M, 1(2):3, 2006. Josh Tobin, Rachel Fong, Alex Ray, Jonas Schneider, Wojciech Zaremba, and Pieter Abbeel. Domain randomization for transferring deep neural networks from simulation to the real world. In 2017 IEEE/RSJ international conference on intelligent robots and systems (IROS), pages 23–30. IEEE, 2017. Misha Laskin, Kimin Lee, Adam Stooke, Lerrel Pinto, Pieter Abbeel, and Aravind Srinivas. Reinforcement learning with augmented data. Advances in neural information processing systems, 33:19884–19895, 2020a.

Denis Yarats, Rob Fergus, Alessandro Lazaric, and Lerrel Pinto. Mastering visual continuous control: Improved data-augmented reinforcement learning. arXiv preprint arXiv:2107.09645, 2021. Taco Cohen and Max Welling. Group equivariant convolutional networks. In International conference on machine learning, pages 2990–2999. PMLR, 2016. Michael M Bronstein, Joan Bruna, Taco Cohen, and Petar Veličković. Geometric deep learning: Grids, groups, graphs, geodesics, and gauges. arXiv preprint arXiv:2104.13478, 2021. Bernt Oksendal. Stochastic differential equations: an introduction with applications. Springer Science & Business Media, 2013. Ioannis Karatzas and Steven Shreve. Brownian motion and stochastic calculus. springer, 2014. Wendell H Fleming and H Mete Soner. Controlled Markov processes and viscosity solutions. Springer, 2006. Jiongmin Yong and Xun Yu Zhou. Stochastic controls: Hamiltonian systems and HJB equations, volume 43. Springer Science & Business Media, 1999. Vivek S Borkar and Vivek S Borkar. Stochastic approximation: a dynamical systems viewpoint, volume 100. Springer, 2008. Harold J Kushner. Numerical methods for stochastic control problems in continuous time. SIAM Journal on Control and Optimization, 28(5):999–1048, 1990. Michael G Crandall, Hitoshi Ishii, and Pierre-Louis Lions. User’s guide to viscosity solutions of second order partial differential equations. Bulletin of the American mathematical society, 27(1):1–67, 1992. Martino Bardi, Italo Capuzzo Dolcetta, et al. Optimal control and viscosity solutions of Hamilton-Jacobi-Bellman equations, volume 12. Springer, 1997. Guy Barles. Solutions de viscosité des équations de Hamilton-Jacobi, volume 17. Springer, 1994. Risi Kondor and Shubhendu Trivedi. On the generalization of equivariance and convolution in neural networks to the action of compact groups. In International conference on machine learning, pages 2747–2755. PMLR, 2018. Maurice Weiler and Gabriele Cesa. General e (2)-equivariant steerable cnns. Advances in neural information processing systems, 32, 2019. Haggai Maron, Heli Ben-Hamu, Nadav Shamir, and Yaron Lipman. Invariant and equivariant graph networks. arXiv preprint arXiv:1812.09902, 2018. Michael J Hutchinson, Charline Le Lan, Sheheryar Zaidi, Emilien Dupont, Yee Whye Teh, and Hyunjik Kim. Lietransformer: Equivariant self-attention for lie groups. In International conference on machine learning, pages 4533–4543. PMLR, 2021. David Ha and Jürgen Schmidhuber. World models. arXiv preprint arXiv:1803.10122, 2(3):440, 2018. Danijar Hafner, Timothy Lillicrap, Jimmy Ba, and Mohammad Norouzi. Dream to control: Learning behaviors by latent imagination. arXiv preprint arXiv:1912.01603, 2019. Michael Laskin, Aravind Srinivas, and Pieter Abbeel. Curl: Contrastive unsupervised representations for reinforcement learning. In International conference on machine learning, pages 5639–5650. PMLR, 2020b. Adam Stooke, Kimin Lee, Pieter Abbeel, and Michael Laskin. Decoupling representation learning from reinforcement learning. In International conference on machine learning, pages 9870–9879. PMLR, 2021. Peter J Olver. Applications of Lie groups to differential equations, volume 107. Springer Science & Business Media, 1993. George W Bluman, Sukeyuki Kumei, and GW Bluman. Symmetries and differential equations, volume 81. Springer, 1989. Nail H Ibragimov. CRC Handbook of Lie Group Analysis of Differential Equations, Volume III. CRC press, 2024. John M Lee. Smooth manifolds. In Introduction to smooth manifolds, pages 1–29. Springer, 2003.

Xue Bin Peng, Marcin Andrychowicz, Wojciech Zaremba, and Pieter Abbeel. Sim-to-real transfer of robotic control with dynamics randomization. In 2018 IEEE international conference on robotics and automation (ICRA), pages 3803–3810. IEEE, 2018. Lerrel Pinto, James Davidson, Rahul Sukthankar, and Abhinav Gupta. Robust adversarial reinforcement learning. In International conference on machine learning, pages 2817–2826. PMLR, 2017. Aravind Rajeswaran, Sarvjeet Ghotra, Balaraman Ravindran, and Sergey Levine. Epopt: Learning robust neural network policies using model ensembles. arXiv preprint arXiv:1610.01283, 2016. Emanuel Todorov, Tom Erez, and Yuval Tassa. Mujoco: A physics engine for model-based control. In 2012 IEEE/RSJ international conference on intelligent robots and systems, pages 5026–5033. IEEE, 2012. Greg Brockman, Vicki Cheung, Ludwig Pettersson, Jonas Schneider, John Schulman, Jie Tang, and Wojciech Zaremba. Openai gym. arXiv preprint arXiv:1606.01540, 2016. Zuyuan Zhang, Mahdi Imani, and Tian Lan. Modeling other players with bayesian beliefs for games with incomplete information. arXiv preprint arXiv:2405.14122, 2024. Zuyuan Zhang, Vaneet Aggarwal, and Tian Lan. Network diffuser for placing-scheduling service function chains with inverse demonstration. In IEEE INFOCOM 2025-IEEE Conference on Computer Communications, pages 1–10. IEEE, 2025b. Zuyuan Zhang and Tian Lan. Lipschitz lifelong monte carlo tree search for mastering non-stationary tasks. arXiv preprint arXiv:2502.00633, 2025. Zuyuan Zhang, Arnob Ghosh, and Tian Lan. Tail-risk-safe monte carlo tree search under pac-level guarantees. arXiv preprint arXiv:2508.05441, 2025c. Zeyu Fang, Zuyuan Zhang, Mahdi Imani, and Tian Lan. Manifold-constrained energy-based transition models for offline reinforcement learning. arXiv preprint arXiv:2602.02900, 2026. Jing Qiao, Zuyuan Zhang, Sheng Yue, Yuan Yuan, Zhipeng Cai, Xiao Zhang, Ju Ren, and Dongxiao Yu. Br-defedrl: Byzantine-robust decentralized federated reinforcement learning with fast convergence and communication efficiency. In IEEE infocom 2024-IEEE conference on computer communications, pages 141–150. IEEE, 2024. Amirhossein Ravari, Guangyu Jiang, Zuyuan Zhang, Mahdi Imani, Robert H Thomson, Aryn A Pyke, Nathaniel D Bastian, and Tian Lan. Adversarial inverse learning of defense policies conditioned on human factor models. In 2024 58th Asilomar Conference on Signals, Systems, and Computers, pages 188–195. IEEE, 2024. Zuyuan Zhang, Vaneet Aggarwal, and Tian Lan. Lisfc-search: Lifelong search for network sfc optimization under non-stationary drifts. arXiv preprint arXiv:2602.14360, 2026d. Sizhe Tang, Jiayu Chen, and Tian Lan. Malinzero: Efficient low-dimensional search for mastering complex multi-agent planning. arXiv preprint arXiv:2511.06142, 2025. Sizhe Tang, Rongqian Chen, and Tian Lan. Agent alpha: Tree search unifying generation, exploration and evaluation for computer-use agents. arXiv preprint arXiv:2602.02995, 2026. Yu Li, Sizhe Tang, Rongqian Chen, Fei Xu Yu, Guangyu Jiang, Mahdi Imani, Nathaniel D Bastian, and Tian Lan. Acdzero: Graph-embedding-based tree search for mastering automated cyber defense. arXiv preprint arXiv:2601.02196, 2026.

Algorithm

A.1

Using learned transformations in RL

Given a transition tuple (s, a, s′ , r) collected from the environment, transition augmentation generates (s̃, ã, s̃′ , r̃) := gα (s), hα (a), gα (s′ ), r , and inserts it into the replay buffer. Under exact value preservation, r̃ = r and the augmented transition is distributed as a valid transition from (s̃, ã). Under approximate value preservation, the 12

Algorithm 1 Value-preserving structure discovery and reinforcement learning with learned transformations 1: Initialize replay buffer D; models (b̂ω , Σ̂ω , r̂ω ); fields (Xθ , Yϕ ); RL parameters ψ 2: for iterations = 1, 2, . . . do 3: Collect rollouts with policy πψ ; store (s, a, s′ , r) into D

Update (b̂ω , Σ̂ω , r̂ω ) on minibatches from D if unknown (and set Q̂ω = Σ̂ω Σ̂⊤ ω) Sample minibatch {(si , ai )} from D Compute residuals Rr , Rb , RQ via automatic differentiation Update (θ, ϕ) by (stochastic) gradient steps on (26) Numerically integrate (27) for small α to obtain (gα , hα ) Use (gα , hα ) for transition augmentation and/or transformation-consistency regularization; update ψ 10: end for 4: 5: 6: 7: 8: 9:

augmented transition is a controlled perturbation whose effect is governed by the generator/reward mismatch bounds in Section 4 and by the discrete-time analogue in Section B. We also encourage the learned critic/value and policy to respect the discovered transformation. For a value function Vψ , we use the invariance penalty h 2 i Rinv (ψ) := Es∼ρS , α∼ν Vψ (s) − Vψ (gα (s)) , where ν is a distribution over small flow times, such as the uniform distribution on [−α0 , α0 ]. For a stochastic policy πψ (· | s), we use a policy-transport consistency penalty based on the pushforward (hα )# πψ (· | s): h i Rpol (ψ) := Es∼ρS , α∼ν KL πψ (· | gα (s)) ∥ (hα )# πψ (· | s) . In implementation, this term can be estimated by sampling a ∼ πψ (· | s) and evaluating the logdensity of hα (a) under πψ (· | gα (s)). The RL parameters are updated with the base actor–critic loss plus these consistency penalties: min LRL (ψ) + λinv Rinv (ψ) + λpol Rpol (ψ). ψ

Theory Guarantees

We provide three layers of guarantees that mirror the pipeline in Section 5. First, we give standard nonconvex optimization guarantees for stochastic structure discovery under smoothness and bounded variance. Second, under realizability and consistent dynamics estimation, we state a consistency result showing that the learned fields converge (in data measure) to a solution of the determining equations, implying vanishing operator mismatch on the data support. Third, we bound numerical errors introduced by finite-step ODE solvers when exponentiating the learned vector fields. Finally, to give an end-to-end non-asymptotic statement free of actor-training nonconvexity, we include a discrete-time tabular analogue showing how exact/approximate value-preserving augmentation perturbs the Bellman fixed point. B.1

Optimization Guarantees for Structure Discovery

Assumption B.1 (Smoothness and bounded variance for structure-discovery optimization). Let Lsym (θ, ϕ) denote the (penalized) objective in (26). Assume Lsym is L-smooth in (θ, ϕ), i.e., its gradient is L-Lipschitz. Assume the stochastic gradients used by the algorithm are unbiased and have uniformly bounded second moments: E∥gt − ∇Lsym (θt , ϕt )∥22 ≤ σ 2 . Theorem B.2 (Convergence of stochastic structure discovery to a stationary point). Under Assumption B.1, running SGD with constant stepsize η ≤ 1/L yields T −1 2 Lsym (θ0 , ϕ0 ) − Linf 1 X 2 E ∥∇Lsym (θt , ϕt )∥2 ≤ + ηLσ 2 , (28) T t=0 ηT 13

where Linf is a lower bound of Lsym . With appropriately √ diminishing stepsizes, one obtains the standard rate min0≤t≤T −1 E∥∇Lsym (θt , ϕt )∥22 = O(1/ T ). This result does not claim global recovery of the true symmetry fields in general (the problem is nonconvex), but it ensures that the optimization procedure reaches approximate stationary points at the usual SGD rates. Consistency requires additional realizability assumptions, addressed next. B.2

Consistency under Realizability

Assumption B.3 (Realizability and consistent dynamics estimation). Assume there exist parameters ω ⋆ , θ⋆ , ϕ⋆ such that b̂ω⋆ = b, Q̂ω⋆ = Q, and r̂ω⋆ = r on the support of ρ, and (Xθ⋆ , Yϕ⋆ ) satisfies the determining equations (12), (13), and (14) on the support of ρ. Assume the estimators (b̂ω , Q̂ω , r̂ω ) are consistent as data grows, and the structure-discovery objective (26) can be optimized to its global minimum. Theorem B.4 (Consistency of learned infinitesimal generators (on data support)). Under Assumption B.3, as the dataset size tends to infinity and the optimization error tends to zero, the learned generators (Xθ , Yϕ ) converge in ρ-mean to a solution of the determining equations (12)–(14). In particular, the induced operator mismatch in Definition 4.1 satisfies εL → 0 and εr → 0 on the support of ρ. Even under realizability, infinitesimal generators are typically identifiable only up to choices of basis within the underlying transformation algebra. The normalization term in (26) fixes the overall scale of Xθ but does not preclude equivalent generators related by linear combinations related by linear combinations when multiple transformation directions are present. B.3

Numerical Flow Approximation

Assumption B.5 (Numerical flow integration). Assume Xθ and Yϕ are Lipschitz on compact sets KS ⊆ S and KA ⊆ A that contain the states/actions encountered by the algorithm. Let ĝα and ĥα be computed by an order-p one-step ODE solver with step size h for α in a bounded interval. Theorem B.6 (Flow approximation error). Under Assumption B.5, for α in a bounded interval, the numerical flows satisfy sup ∥ĝα (s) − gα (s)∥2 ≤ Cg hp ,

sup ∥ĥα (a) − hα (a)∥2 ≤ Ch hp ,

s∈KS

a∈KA

(29)

where Cg and Ch depend on Lipschitz constants of the fields on KS , KA and on the solver stability region. B.4

A Discrete-Time Analogue: Value-Preserving Augmented Bellman Backups

The continuous-time theory above explains why (approximate) generator equivariance leads to (approximate) value invariance. To provide a fully non-asymptotic end-to-end convergence statement that avoids nonconvex actor training, we additionally state a discrete-time tabular analogue in which value iteration is a contraction and fixed-point perturbations can be bounded explicitly. Let T be the optimal Bellman operator in a discounted discrete-time MDP: Z (T V )(s) = max r(s, a) + γ V (s′ ) P (ds′ | s, a) . a∈A

(30)

T is a γ-contraction in ∥ · ∥∞ . Theorem B.7 (Exact value-preserving augmentation preserves the Bellman operator). Assume there exists a transformation pair (g, h) forming an exact value-preserving correspondence for the discrete-time MDP, namely r(g(s), h(a)) = r(s, a) and P (· | g(s), h(a)) = g# P (· | s, a) for all (s, a). Then the Bellman backup is invariant under augmentation by (g, h) in expectation, so tabular value iteration with exact value-preserving augmentation converges to V ⋆ at the usual rate ∥Vk − V ⋆ ∥∞ ≤ γ k ∥V0 − V ⋆ ∥∞ . 14

Theorem B.8 (Approximate value-preserving augmentation R yields a bounded fixed-point perturbation). Define total variation as dTV (µ, ν) := sup∥f ∥∞ ≤1 f d(µ − ν) . Assume the augmentation is (εP , εr ) approximate in the discrete-time sense: sup r(g(s), h(a)) − r(s, a) ≤ εr , sup dTV P (· | g(s), h(a)), g# P (· | s, a) ≤ εP . (31) s,a

s,a

Let T̃ be the Bellman operator induced by the approximately augmented model. Then for any bounded V, ∥(T̃ V ) − (T V )∥∞ ≤ εr + γ εP ∥V ∥∞ , (32) and T̃ remains a γ-contraction. Let Ṽ ⋆ be the fixed point of T̃ . Then the fixed-point error satisfies ∥Ṽ ⋆ − V ⋆ ∥∞ ≤ B.5

1 (εr + γ εP ∥V ⋆ ∥∞ ) . 1−γ

(33)

HJB well-posedness used in the main text

Assumption B.9 (HJB well-posedness). Assume: (i) A is a compact metric space and r, b, Σ are continuous in (s, a) and Lipschitz in s uniformly over a; moreover |r| ≤ Rmax and b, Σ are bounded. (ii) The diffusion matrix Q(s, a) = Σ(s, a)Σ(s, a)⊤ is uniformly elliptic: there exists λ > 0 such that ξ ⊤ Q(s, a)ξ ≥ λ∥ξ∥2 for all (s, a) and ξ ∈ Rd . (iii) Either S is bounded with an appropriate state-constraint or boundary condition, or S = Rd with a growth/boundedness condition ensuring that V ⋆ is well defined and bounded. (iv) A comparison principle holds for bounded viscosity suband supersolutions of (7).

Running Example: Planar Rotation as a Special Case

0 Let ê3 = 1 0 "

−1 0 0

# 0 0 , 0

gθ = exp(θê3 ) ∈ SO(3), and let the induced action on vectors v ∈ R3 cos θ Rz (θ) = sin θ 0 "

be ρ(gθ )v = Rz (θ)v,

− sin θ cos θ 0

# 0 0 , 1

θ ∈ (−π, π].

Consider a state s = (x, xd , ξ), where x ∈ R3 is the current position, xd ∈ R3 is a goal/location descriptor, and ξ collects scalar quantities left invariant by a transformation family. Let the action be a control vector a = F ∈ R3 . For the planar-orbiting case, define gθ · s := Rz (θ)x, Rz (θ)xd , ξ , hθ · a := Rz (θ)a. If the controlled dynamics and reward depend only on relative geometry and transformation-invariant control quantities, then the generator-level value-preserving condition and reward compatibility hold: La Ugθ = Ugθ Lhθ (a) , r(gθ·s, hθ·a) = r(s, a). Hence the optimal policy can be transported consistently along the orbit, π ⋆ (gθ·s) ≈ hθ·π ⋆ (s). The same perspective extends beyond globally symmetric coordinates. For a UAV executing a figureeight maneuver, one can align repeated phase-matched relative configurations across the two lobes of the trajectory. These matched configurations need not form a global coordinate symmetry, but they can still induce approximate value-preserving correspondences when the local control semantics, geometry, and reward-to-go remain close. This illustrates why the target of our method is broader value-preserving structure rather than symmetry alone.

Additional Experimental Details and Full Results

This appendix provides full environment specifications, hyperparameters, and complete plots for all benchmark variants. We also include ablations and additional diagnostics linking generator discovery to operator mismatch and value invariance. Unless noted otherwise, all environments are implemented in gymnasium style with continuous observation and action spaces, and at submission time, the anonymized supplementary material contains key implementation code for VPSD-RL, configuration templates. The complete cleaned repository and raw logs will be released upon acceptance. 15

𝜋 ∗ (𝑠)

𝐹

𝑥!

𝜋 ∗ (𝑠)

Rotational Symmetry:

𝐹

cos 𝜃 𝜌 𝑔# = 𝑅$ 𝜃 = sin 𝜃 0

𝑥!

𝑎∗

𝑆

𝑆 𝑔#

𝜋 ∗ (𝑠̃ )

𝑆*

𝑆* = 𝑔# 𝑆

Standard RL / no discovered structure

𝑎∗

True Symmetry:

𝑔#

0 −1 0 𝑔# = 𝑒𝑥𝑝 𝜃𝑒̂% , 𝑒̂% = 1 0 0 0 0 0

𝜋 ∗ (𝑠̃ )

−sin 𝜃 cos 𝜃 0

𝑔# 𝐹

0 0 1

Learned Symmetry: 0 −1 0 𝑔># = 𝑒𝑥𝑝 𝜃𝐴@ , 𝐴@ = 1 0 0 0 0 0

𝑎-∗ = 𝑔# 𝑎∗ Top: optimization curve up to step 400. Bottom: 𝑒̂% and learned matrices at selected stages.

True symmetry and learned symmetry

Figure 3: Running planar-rotation example illustrating the role of discovered value-preserving structure in continuous control. Left: a standard policy learner has no explicit mechanism tying a visited state s to its transformed counterpart s̃ = gθ s. Middle: in the rotational special case, the ground-truth transformation is generated by gθ = exp(θê3 ), with induced action ρ(gθ )v = Rz (θ)v; our method instead learns an infinitesimal generator Â and constructs the finite transform ĝθ = exp(θÂ). Right: numerical diagnostics show that, as optimization proceeds, the learned generator approaches the ground-truth rotational generator and the induced finite transforms become increasingly consistent with true orbit transport, supporting the value-preserving relation π ⋆ (gθ s) ≈ hθ π ⋆ (s). D.1

Stage I: Full synthetic environments and diagnostics

Common state/action conventions. All three synthetic environments use the same state and action parameterization. The state is s = (x, y, vx , vy ) ∈ R4 (position and velocity in the plane), and the action is a bounded acceleration command a = (ax , ay ) ∈ [−1, 1]2 . Let p = (x, y) and v = (vx , vy ). All stochasticity enters through the velocity channels only. Each environment is an Euler–Maruyama discretization of a controlled diffusion of the form dpt = vt dt,

dvt = f (pt , vt , at ) dt + σ dWt ,

(34)

where Wt is a 2D standard Wiener process and σ > 0 is isotropic. In discrete time with step size ∆t we use the semi-implicit update √ vk+1 = vk + ∆t f (pk , vk , ak ) + σ ∆t ξk , pk+1 = pk + ∆t vk+1 , (35) with ξk ∼ N (0, I2 ) i.i.d. Rot2D-Point (Rot2DPointEnv). This is the baseline SO(2)-invariant controlled diffusion with linear damping: f (p, v, a) = a − λv, (36) where λ > 0 is the damping coefficient. The instantaneous reward is SO(2)-invariant: r(s, a) = −∥p∥22 − 0.1 ∥v∥22 .

(37)

The induced continuous symmetry action is the usual planar rotation: for any R ∈ SO(2), gR (p, v, a) = (Rp, Rv, Ra),

(38)

under which both the drift (36) and reward (37) are invariant. The environment uses defaults ∆t = 0.05, σ = 0.02, λ = 0.1, horizon T = 200 steps (see code listing). Ref2D-DoubleWell (Ref2DDoubleWellEnv). This environment uses a double-well potential with reflection symmetry x 7→ −x when δ = 0. Define the base potential U0 (x, y) = (x2 − 1)2 + y 2 ,

(39)

and the dynamics potential (used in the drift) with a tunable symmetry-breaking term Udyn (x, y) = U0 (x, y) + δx.

(40)

f (p, v, a) = −∇Udyn (p) − λv + a,

(41)

The drift is

Table 2: Synthetic environment default parameters in our implementation. Env

∆t

Horizon

Extra

Notes

Rot2D-Point Ref2D-DoubleWell PostConstraint-Rot2D

0.05 0.05 0.05

0.02 0.02 0.02

0.1 0.1 0.1

200 200 200

– δ [rmin , rmax ]

SO(2) invariant reflection exact if δ = 0 projection + tangential velocity

where ∇Udyn (x, y) = 4x(x2 − 1) + δ, 2y . The reward is kept independent of δ (matching the implementation): r(s, a) = −U0 (x, y) − 0.1 ∥v∥22 . (42) When δ = 0, the environment is invariant under the reflection action g(p, v, a) = ((−x, y), (−vx , vy ), (−ax , ay )).

(43)

When δ ̸= 0, the symmetry is only approximate; we use this setting to stress-test discovery under controlled mismatch. The implementation additionally applies mild clipping to position/velocity to avoid numerical blow-ups (see code listing). Defaults: ∆t = 0.05, σ = 0.02, λ = 0.1, horizon T = 200 steps. PostConstraint-Rot2D (PostConstraintRot2DEnv). This environment starts from the Rot2D drift (36) and then enforces a feasibility constraint on position: ∥p∥2 ∈ [rmin , rmax ].

(44)

After the unconstrained update (35), we project the position to the annulus: pk+1 ← ΠA (pk+1 ),

A = {p : rmin ≤ ∥p∥2 ≤ rmax },

(45)

and when a projection occurs we remove the radial velocity component to maintain boundary consistency: pk+1 vk+1 ← vk+1 − u ⟨vk+1 , u⟩, u= . (46) ∥pk+1 ∥2 The reward matches Rot2D-Point: r(s, a) = −∥p∥22 − 0.1 ∥v∥22 .

(47)

This environment preserves rotational symmetry in the interior but introduces non-smooth geometry and boundary effects, making it a convenient testbed for constrained flows and retraction/projection handling in the discovered transformations. Defaults: ∆t = 0.05, σ = 0.02, λ = 0.1, rmin = 0.5, rmax = 2.0, horizon T = 200 steps. The step function also logs whether projection occurred (info["projected"]). Summary of synthetic parameters. Table 2 lists the default parameters used by our reference implementation (exact values in the released config files). Symmetry discovery setup. For Stage I, we collect transition tuples (s, a, r, s′ ) by rolling out a behavior policy (Appendix D.3). The discovery module fits a generator field (or parameterization thereof) by minimizing determining-equation residuals. We report: (i) residual statistics E∥Rb ∥22 , E∥RQ ∥2F , E|Rr |2 ; (ii) generator alignment with ground truth (cosine similarity after normalization); and (iii) value non-invariance Es,α |V (s) − V (gα (s))| vs. |α|. Additional plots. D.2

Stage II: SymNav-15 benchmark variants

State, action, and observation. SymNavEnv is a 2D navigation environment with obstacles and smooth wind disturbances. The latent physical state is again s = (x, y, vx , vy ) ∈ R4 and the action is a = (ax , ay ) ∈ [−1, 1]2 . The observation augments state with goal-relative features and lidar rays: o = x, y, vx , vy , (gx − x), (gy − y), d1 , . . . , d8 ∈ R14 , (48) where di ∈ [0, 1] is the normalized distance-to-collision along ray i (8 rays uniformly spaced over [0, 2π)), computed against both circular obstacles and the world boundary (a disk). 17

60K

120K

180K

Env steps

(a)

240K

300K

PostConstraint-Rot2D-v0 Ref2D-DoubleWell-v0 Rot2D-Point-v0

60K

120K

180K

Env steps

240K

300K

(b)

0.50 0.45 0.40 0.35 0.30 0.25 0.20 0.15 0.10

Regularizer-to-RL loss ratio

Determining-equation residual loss PostConstraint-Rot2D-v0 Ref2D-DoubleWell-v0 Rot2D-Point-v0

)

Structure regularization (symmetry consistency)

sym/( q +

0.50 0.45 0.40 0.35 0.30 0.25 0.20 0.15 0.10

loss_eq

PostConstraint-Rot2D-v0 Ref2D-DoubleWell-v0 Rot2D-Point-v0

loss_sym

loss_inv

Invariance violation proxy (training batch-averaged) 1.75 1.50 1.25 1.00 0.75 0.50 0.25 0.00

60K

120K

180K

Env steps

240K

300K

2 1 0 1 2 3 4 5

PostConstraint-Rot2D-v0 Ref2D-DoubleWell-v0 Rot2D-Point-v0

60K

(c)

120K

180K

Env steps

240K

300K

(d)

Figure 4: Fig. 4a Invariance violation proxy over training (batch-averaged Linv ). Fig. 4b Symmetry/consistency regularization term over training (batch-averaged Lsym ). Fig. 4c Determiningequation residual term used during Stage I (batch-averaged Leq ). Fig. 4d Regularizer-to-RL loss ratio across tasks, illustrating that regularization does not dominate optimization. Dynamics. The environment implements the same semi-implicit update (35) with a variantdependent wind field: f (p, v, a) = a + w(p) − λv, (49) where λ > 0 is damping and w(p) is a smooth wind: w(x, y) = k sin(0.5 y), cos(0.5 x) , k = 0.15 (variant mod 5) + 1 . (50) Noise is isotropic in the velocity channels, as in (34). After each step, position is clamped to remain inside the world disk of radius Rworld by radial rescaling. Map generation (variants). Each variant is identified by an integer variant ∈ {1, . . . , 15}. The goal is placed on a circle of radius rg with angle determined by the variant index: (variant mod 15) , rg = 3.5. (51) 15 Obstacles are 6 circles with centers and radii deterministically generated from variant (matching the implementation): for k ∈ {0, . . . , 5}, 2πk θk = + 0.15(variant mod 5), ρk = 1.2 + 0.2 (variant + k) mod 3 , (52) 6 ck = ρk (cos θk , sin θk ), rk = 0.35 + 0.05 (variant + 2k) mod 4 . (53) This generator yields 15 related maps with shared observation/action interfaces but diverse local geometry. g = rg (cos θ, sin θ),

θ = 2π

Reward, termination, and diagnostics. Let d(p, g) = ∥p − g∥2 be the goal distance. The step reward is dense shaping minus collision penalty plus a terminal success bonus: r = −d(p, g) − 2 · I{collision} + 10 · I{d(p, g) < ε},

ε = 0.3.

(54)

Episodes terminate on success (d(p, g) < ε) and truncate at horizon T = 400. The environment returns diagnostic flags info["success"], info["collided"], and info["dist"]. Approximate transformations. While the exact rotational symmetry is broken by variant-specific obstacles, goal placement, and wind scaling, the construction is designed to admit approximate continuous transformations: rotating (p, v, a) and rotating the map/goal yields near-matching local dynamics in many regions, with mismatch localized near obstacles and due to non-invariant wind terms. This is precisely the regime where operator mismatch bounds and value non-invariance diagnostics are most informative. Variant list and full learning curves. We evaluate variants {V1, . . . , V15} corresponding to variant=1..15. Full learning curves for all variants and all methods are in Fig. 5. (We optionally include a compact table of map/goal/wind parameters per variant in the released appendix PDF.) Training details. All methods use the same observation preprocessing, episode truncation, and evaluation protocol. Per-method hyperparameters (optimizer, learning rate, batch size, replay settings for off-policy methods, etc.) and the exact seed list are provided in Appendix D.3 and the accompanying configuration files. 18

SymNav-v1

VPSD-RL SAC TD3 PPO DDPG A2C DynRand RARL EPOPT

1000 1250 1500 1750 0

60K

120K

180K

Env steps

500

750 VPSD-RL SAC TD3 PPO DDPG A2C DynRand RARL EPOPT

1000 1250 1500 1750

240K

2000

300K

(a) SymNav-V1

1250 1750 0

60K

120K

180K

Env steps

240K

VPSD-RL SAC TD3 PPO DDPG A2C DynRand RARL EPOPT

1000 1250

120K

180K

Env steps

240K

2000

1250 1500

VPSD-RL SAC TD3 PPO DDPG A2C DynRand RARL EPOPT

750 1000 1250

2000

300K

500

750

VPSD-RL SAC TD3 PPO DDPG A2C DynRand RARL EPOPT

1000 1250 1500 1750 60K

120K

180K

Env steps

240K

Eval return

250

750

120K

180K

Env steps

240K

1250

2000

300K

VPSD-RL SAC TD3 PPO DDPG A2C DynRand RARL EPOPT

1000 1250

2000

1250 1500

60K

120K

180K

Env steps

240K

2000

300K

Env steps

(m) SymNav-V13

VPSD-RL SAC TD3 PPO DDPG A2C DynRand RARL EPOPT

1000 1250 1500 1750

240K

300K

2000

60K

120K

180K

Env steps

240K

Eval return

500

180K

120K

180K

Env steps

240K

300K

SymNav-v15

750

120K

60K

(l) SymNav-V12

SymNav-v14

500

60K

(k) SymNav-V11

750

VPSD-RL SAC TD3 PPO DDPG A2C DynRand RARL EPOPT

1000

500

2000

300K

SymNav-v12

750

250

1750

240K

500

250

1500

180K

Env steps

1750 0

VPSD-RL SAC TD3 PPO DDPG A2C DynRand RARL EPOPT

120K

250

1750

300K

60K

(i) SymNav-V9

250

1250

1500

SymNav-v13

1000

VPSD-RL SAC TD3 PPO DDPG A2C DynRand RARL EPOPT

1000 1500

SymNav-v11

(j) SymNav-V10 0

60K

250

300K

SymNav-v9

(h) SymNav-V8

SymNav-v10

240K

1750 0

(g) SymNav-V7 0

180K

Env steps

500

Eval return

Env steps

240K

120K

250

1750 180K

60K

(f) SymNav-V6

1500

1750 120K

Eval return

1000

60K

VPSD-RL SAC TD3 PPO DDPG A2C DynRand RARL EPOPT

1250 1750

300K

SymNav-v8

500

Eval return

VPSD-RL SAC TD3 PPO DDPG A2C DynRand RARL EPOPT

300K

750 1000 1500

250

750

240K

SymNav-v6

(e) SymNav-V5

500

Eval return

60K

250

180K

Env steps

500

750

2000

120K

1750

300K

60K

1500

SymNav-v7

250

(d) SymNav-V4 0

Eval return

2000

Eval return

1000

Eval return

VPSD-RL SAC TD3 PPO DDPG A2C DynRand RARL EPOPT

1500

Eval return

1750

300K

500

750

2000

240K

250

500

2000

180K

Env steps

VPSD-RL SAC TD3 PPO DDPG A2C DynRand RARL EPOPT

1250 1500

SymNav-v5

250

2000

120K

750 1000

(b) SymNav-V2

SymNav-v4

60K

Eval return

250

500

750

SymNav-v3

250

500

2000

SymNav-v2

250

300K

(n) SymNav-V14

750 VPSD-RL SAC TD3 PPO DDPG A2C DynRand RARL EPOPT

1000 1250 1500 1750 2000

60K

120K

180K

Env steps

240K

300K

(o) SymNav-V15

Figure 5: Full learning curves on all SymNav-15 variants for all compared methods.

D.3

Implementation details: model fitting and flow integration

Data collection for discovery. Discovery operates on transition tuples (s, a, r, s′ ) collected online during RL training. For on-policy methods, we use recent rollouts; for off-policy methods, we sample from the replay buffer. To stabilize the determining-residual optimization, we (i) balance minibatches across the current state distribution, and (ii) normalize state components (position/velocity scales) using running statistics. Dynamics and reward modeling. For synthetic environments, ground-truth drift/diffusion/reward are known analytically via Eqs. (36)–(47), (41)–(42), and are used to compute ground-truth generators 19

Table 3: Aggregate results on SymNav-V1. Numbers are mean ± std over seeds, averaged across all variants. AUC is computed up to a fixed step budget. Wall-clock reports seconds per 105 environment steps. Method

Avg. normalized return ↑

AUC ↑

Wall-clock (sec / 105 steps) ↓

A2C DDPG PPO SAC TD3

-1876.58±321.57 -106.97±2.36 -270.43±136.65 -106.60±1.97 -104.58±0.41

-1952.71±217.57 -206.47±39.70 -424.68±140.79 -233.54±43.70 -172.08±25.35

175±1 540±7 217±35 410±12 754±10

DynRand RARL EPOpt

-115.54±7.25 -114.53±6.48 -197.45±67.57

-181.49±18.78 -155.27±17.23 -271.25±78.60

581±12 753±12 485±27

VPSD-RL (Aug+Reg)

-124.48±12.65

-145.78±24.65

752±20

and diagnostic alignment. For complex environments (SymNav and MuJoCo), the discovery module fits local surrogate models for drift b(s, a) and reward r(s, a) by regression on one-step transitions: ′

bb(s, a) ≈ s − s , ∆t

rb(s, a) ≈ r.

If a diffusion model is used, we parameterize Q(s, a) through a PSD form (e.g., diagonal or Cholesky factor) to ensure positive semidefiniteness. All exact architectural choices and optimizer settings are provided in the config files. Flow integration and constraint handling. Given a learned generator X(·) (and Y (·) when needed), we produce finite transformations by integrating the ODE d z(α) = X(z(α)), z(0) = z0 , dα (or its coupled version in the method section), using a fixed-step explicit solver (RK4 by default). We integrate over a symmetric range α ∈ [−αmax , αmax ] with step size h. For constrained environments (PostConstraint-Rot2D), we apply the same projection/retraction policy as the environment: position is projected to the feasible annulus and the radial velocity component is removed when projection occurs (Eqs. (45)–(46)). We reject transformations that cause numerical instability (NaNs/Infs) and fall back to identity augmentation in that case.

Limitations and broader impact.

VPSD-RL relies on differentiable surrogate drift, diffusion, and reward estimates, as well as locally valid Lie-group flows. Its performance may degrade when the learned generator is poorly identified, when the dynamics are strongly discontinuous, or when the value-preserving relation holds only on a narrow data support. The current evaluation is limited to controlled simulation benchmarks, so deployment in safety-critical robotic or physical systems requires additional validation, constraint checking, and monitoring under distribution shift. The expected positive impact is improved data efficiency and robustness in continuous-control RL; a possible negative impact is over-trusting learned transformations in high-stakes autonomous systems when the approximate-invariance assumptions fail.

Proof

F.1

Proof of Theorem 2.1

Proof. Fix β > 0 and consider the controlled diffusion dst = b(st , at ) dt + Σ(s , at ) dWt with at ∈ Rt∞ A admissible. For any Markov policy π, the discounted value is V π (s) = Eπ 0 e−βt r(st , at ) dt and V ⋆ (s) = supπ V π (s). Since |r| ≤ Rmax , for every π and every s, Z ∞ Z ∞ Rmax |V π (s)| ≤ Eπ e−βt |r(st , at )| dt ≤ e−βt Rmax dt = , β 0 0 20

hence V ⋆ is well-defined and bounded. Under Assumption B.9(i)–(iii), the SDE is well-posed for admissible controls and the strong Markov property together with control concatenation yields the discounted dynamic programming principle: for any bounded stopping time τ , hZ τ i V ⋆ (s) = sup Eπ e−βt r(st , at ) dt + e−βτ V ⋆ (sτ ) s0 = s . (55) π

We now verify that V ⋆ is a bounded viscosity solution of the HJB equation βV = (HV ), where for φ ∈ C 2 (S), 1 (La φ)(s) := b(s, a) · ∇φ(s) + tr Q(s, a)∇2 φ(s) , (Hφ)(s) := sup {r(s, a) + (La φ)(s)}. 2 a∈A Let φ ∈ C 2 (S) and suppose V ⋆ − φ attains a local maximum at s0 with V ⋆ (s0 ) = φ(s0 ). Then there exists ρ > 0 such that V ⋆ (s) ≤ φ(s) for all s ∈ Bρ (s0 ). Let τρ := inf{t ≥ 0 : st ∈ / Bρ (s0 )} and fix t > 0; set τ := t ∧ τρ . By (55) and the definition of supremum, for any ε > 0 there exists an admissible policy π ε such that Z τ ⋆ πε −βu −βτ ⋆ V (s0 ) ≤ E e r(su , au ) du + e V (sτ ) + εt. 0 ⋆

Since sτ ∈ Bρ (s0 ), we have V (sτ ) ≤ φ(sτ ), hence using V ⋆ (s0 ) = φ(s0 ), Z τ πε −βu −βτ 0≤E e r(su , au ) du + e φ(sτ ) − φ(s0 ) + εt.

(56)

Apply Itô’s formula to e−βu φ(su ) under π ε . Using dsu = b(su , au ) du + Σ(su , au ) dWu and dsu ds⊤ u = Q(su , au ) du, we obtain 1 dφ(su ) = ∇φ(su ) · b(su , au ) du + ∇φ(su ) · Σ(su , au ) dWu + tr Q(su , au )∇2 φ(su ) du, 2 and therefore d e−βu φ(su ) = e−βu (Lau φ)(su ) − βφ(su ) du + e−βu ∇φ(su ) · Σ(su , au ) dWu . Integrating from 0 to τ gives the identity Z τ Z τ −βτ −βu e φ(sτ ) − φ(s0 ) = e (Lau φ)(su ) − βφ(su ) du + e−βu ∇φ(su ) · Σ(su , au ) dWu . 0

The stochastic integral is a martingale with zero expectation (bounded coefficients and τ ≤ t < ∞), hence Z τ π ε −βτ πε −βu E e φ(sτ ) − φ(s0 ) = E e (Lau φ)(su ) − βφ(su ) du . 0

Substituting into (56) yields Z τ ε 0 ≤ Eπ e−βu r(su , au ) + (Lau φ)(su ) − βφ(su ) du + εt. 0

Define Ψ(s) := supa∈A {r(s, a) + (La φ)(s) − βφ(s)} = (Hφ)(s) − βφ(s). Since r(su , au ) + (Lau φ)(su ) − βφ(su ) ≤ Ψ(su ), Z τ πε −βu 0≤E e Ψ(su ) du + εt. 0

Dividing by t > 0 gives

Z τ 1 πε −βu 0≤ E e Ψ(su ) du + ε. t 0 ε R τ We claim 1t Eπ 0 e−βu Ψ(su ) du → Ψ(s0 ) as t ↓ 0. Indeed, since τ ≤ t and Ψ is bounded on Bρ (s0 ), using |e−βu − 1| ≤ βu we have Z τ Z τ 1 1 β∥Ψ∥∞ (e−βu − 1)Ψ(su ) du ≤ E βu ∥Ψ∥∞ du ≤ E t → 0. t t 2 0 0 21

Moreover, Ψ is uniformly continuous on Bρ (s0 ), so for its modulus of continuity ω(·), the term R τ 1 E vanishes as t ↓ 0 because sup0≤u≤t ∥su − s0 ∥ → 0 in probability |Ψ(s ) − Ψ(s )| du u 0 t 0 √ under bounded coefficients (drift is O(t) and the martingale part is O( t) by BDG), and τ = t ∧ τρ satisfies E[τ ]/t → 1. Consequently, Z τ 1 πε −βu lim E e Ψ(su ) du = Ψ(s0 ). t↓0 t 0 Letting t ↓ 0 yields 0 ≤ Ψ(s0 ) + ε, and then ε ↓ 0 gives Ψ(s0 ) ≥ 0, i.e. βφ(s0 ) ≤ (Hφ)(s0 ). This proves that V ⋆ is a viscosity subsolution. Now suppose V ⋆ − φ attains a local minimum at s0 with V ⋆ (s0 ) = φ(s0 ). Then for some ρ > 0, V ⋆ ≥ φ on Bρ (s0 ), and with τ = t ∧ τρ we have, for any fixed a ∈ A, by (55) (since supπ dominates the constant control au ≡ a), Z τ Z τ ⋆ −βu −βτ ⋆ −βu −βτ V (s0 ) ≥ E e r(su , a) du + e V (sτ ) ≥ E e r(su , a) du + e φ(sτ ) . 0

⋆

Using V (s0 ) = φ(s0 ) and the same Itô identity as above (with generator La ) gives Z τ 0≥E e−βu r(su , a) + (La φ)(su ) − βφ(su ) du . 0

Divide by t and let t ↓ 0 to obtain 0 ≥ r(s0 , a) + (La φ)(s0 ) − βφ(s0 ) for every a ∈ A. Taking the supremum over a yields βφ(s0 ) ≥ (Hφ)(s0 ), hence V ⋆ is a viscosity supersolution. Therefore V ⋆ is a bounded viscosity solution of βV = HV , establishing existence of a bounded viscosity solution. Uniqueness among bounded viscosity solutions follows directly from the comparison principle in Assumption B.9(iv): if u is a bounded viscosity subsolution and v a bounded viscosity supersolution, then u ≤ v. In particular, if V is any bounded viscosity solution, then V is both a subsolution and a supersolution, so V ≤ V ⋆ and V ⋆ ≤ V , implying V ≡ V ⋆ . F.2

Proof of Lemma 3.2

Proof. We fix a ∈ A and write ã := hϑ (a), Tt := Pta , T̃t := Ptã , A := La , Ã := Lã , and U := Ugϑ . We will use two standard C0 -semigroup facts, stated and proved here for completeness. Lemma F.1. Let (St )t≥0 be a strongly continuous semigroup on a Banach space X with generator d G. For every x ∈ Dom(G), the map t 7→ St x is norm-differentiable and satisfies dt St x = GSt x = St Gx for all t ≥ 0. Moreover, if w : [0, ∞) → X is norm-differentiable, w(t) ∈ Dom(G) for all t, and w′ (t) = Gw(t) with w(0) = x ∈ Dom(G), then necessarily w(t) = St x for all t ≥ 0. Proof. Fix x ∈ Dom(G). By definition of G, the limit Gx = limh↓0 Sh x−x exists in X. For t ≥ 0 h and h > 0, use the semigroup property to compute St+h x − St x St (Sh x − x) Sh x − x = = St . h h h tx Since St is bounded for fixed t, letting h ↓ 0 yields the right-derivative limh↓0 St+h x−S = St Gx. h For 0 < h < t, similarly St x − St−h x St−h (Sh x − x) Sh x − x = = St−h , h h h

and letting h ↓ 0 gives the left-derivative equal to St Gx by strong continuity of St−h at t. Thus d t 7→ St x is norm-differentiable and dt St x = St Gx. To identify GSt x, note that for h > 0, Sh (St x) − St x St (Sh x − x) Sh x − x = = St → St Gx (h ↓ 0), h h h d so St x ∈ Dom(G) and GSt x = St Gx, hence also dt St x = GSt x.

For uniqueness, let w satisfy the stated assumptions and fix t > 0. Define ϕ(s) := St−s w(s) for d s ∈ [0, t]. Using the already proved identity ds St−s y = −GSt−s y for y ∈ Dom(G) and the product rule in X, we compute ϕ′ (s) = −GSt−s w(s) + St−s w′ (s) = −GSt−s w(s) + St−s Gw(s) = 0, so ϕ is constant and ϕ(t) = ϕ(0) gives w(t) = St w(0) = St x. Since t was arbitrary, w(t) = St x for all t ≥ 0. Lemma F.2. If (St )t≥0 is a strongly continuous semigroup on a Banach space X with generator G, then Dom(G) is dense in X. Proof. Let λ > 0 and define the bounded operator (the Yosida approximation) Z ∞ Gλ x := λ e−λt St x dt − x, x ∈ X, 0

where the Bochner R ∞integral exists because t 7→ St x is continuous and ∥St ∥ is locally bounded. Define xλ := λ 0 e−λt St x dt. Then xλ → x in X as λ → ∞ by strong continuity at 0 (an Abel-type limit). Moreover, one checks directly from the generator definition that xλ ∈ Dom(G) and Gxλ = λ(xλ − x) (equivalently xλ ∈ Ran(λI − G) ⊂ Dom(G)). Hence every x ∈ X is a norm-limit of elements in Dom(G), proving density. We now prove the two directions of Lemma 3.2. Assume first the generator intertwining holds on the common generator core: ∀f ∈ Dom(Ã).

AU f = U Ãf,

Fix f ∈ Dom(Ã) and define u(t) := U T̃t f for t ≥ 0. By Lemma F.1 applied to (T̃t ) and Ã, we have T̃t f ∈ Dom(Ã) and d T̃t f = ÃT̃t f. dt Since U is bounded linear, u′ (t) = U ÃT̃t f. Because AU = U Ã on Dom(Ã), we have u′ (t) = AU T̃t f = Au(t),

u(0) = U f.

Define v(t) := Tt U f . By Lemma F.1 applied to (Tt ) and A, v ′ (t) = Av(t) and v(0) = U f . Hence u(t) = v(t), namely U T̃t f = Tt U f, ∀t ≥ 0, ∀f ∈ Dom(Ã). By density of Dom(Ã) and boundedness of U, Tt , T̃t , this extends to all f ∈ F: Tt U = U T̃t . Equivalently, h (a)

Pta Ugϑ = Ugϑ Pt ϑ

Conversely, assume the semigroup intertwining holds: Tt U = U T̃t ,

t ≥ 0.

For every f ∈ Dom(Ã), Tt U f − U f U T̃t f − U f = lim = U Ãf. t↓0 t↓0 t t

AU f = lim

Therefore AU = U Ã, equivalently La Ugϑ = Ugϑ Lhϑ (a) , which is Eq. (8) on the generator domain. 23

F.3

Proof of Theorem 3.3

Proof. Fix (s, a) ∈ S × A and abbreviate gϵ := gϑϵ , hϵ := hϑϵ . Reward invariance in Definition 3.1 reads r(gϵ (s), hϵ (a)) = r(s, a) for all ϵ. Differentiating at ϵ = 0 and using the chain rule gives 0=

d d d = ∇s r(s, a) · + ∇a r(s, a) · , r(gϵ (s), hϵ (a)) gϵ (s) hϵ (a) dϵ dϵ dϵ ϵ=0 ϵ=0 ϵ=0

which is exactly (12). For each fixed ϵ, Definition 3.1 gives that for all f ∈ C 2 (S), La (f ◦ gϵ ) (s) = (Lhϵ (a) f ) ◦ gϵ (s) = Lhϵ (a) f (gϵ (s)).

(57)

Write the controlled generator in coordinates as 1 tr Q(x, a) ∇2 f (x) , Q(x, a) = Σ(x, a)Σ(x, a)⊤ . 2 Set xϵ := gϵ (s) and denote the Jacobian matrix and second derivatives of gϵ by Jϵ (s) := ∇s gϵ (s) ∈ Rd×d , ∇2s gϵk (s) ij := ∂ij gϵk (s) , k = 1, . . . , d. (La f )(x) = b(x, a) · ∇f (x) +

For later use we expand the first and second spatial derivatives of f ◦ gϵ at s. By the chain rule, for each i = 1, . . . , d, ∂i (f ◦ gϵ )(s) =

d X

(∂k f )(xϵ ) ∂i (gϵk )(s),

i.e.

∇(f ◦ gϵ )(s) = Jϵ (s)⊤ ∇f (xϵ ).

k=1

Differentiating once more and using product/chain rules yields, for each i, j, ∂ij (f ◦ gϵ )(s) =

d X

(∂kℓ f )(xϵ ) ∂i (gϵk )(s) ∂j (gϵℓ )(s) +

k,ℓ=1

d X

(∂k f )(xϵ ) ∂ij (gϵk )(s).

k=1

Equivalently in matrix form, ∇2 (f ◦ gϵ )(s) = Jϵ (s)⊤ ∇2 f (xϵ )Jϵ (s) +

d X

∂k f (xϵ ) ∇2s gϵk (s).

k=1

Plugging these into the left-hand side of (57) gives 1 La (f ◦ gϵ ) (s) = b(s, a) · ∇(f ◦ gϵ )(s) + tr Q(s, a) ∇2 (f ◦ gϵ )(s) 2 1 ⊤ = b(s, a) · Jϵ (s) ∇f (xϵ ) + tr Q(s, a) Jϵ (s)⊤ ∇2 f (xϵ )Jϵ (s) 2 d X 1 + ∂k f (xϵ ) tr Q(s, a) ∇2s gϵk (s) . 2 k=1

Using cyclicity of trace, define e ϵ (s, a) := Jϵ (s) Q(s, a) Jϵ (s)⊤ , Q and define the correction vector cϵ (s, a) ∈ Rd by cϵ (s, a) k := tr Q(s, a) ∇2s gϵk (s) ,

k = 1, . . . , d.

Then 1 1 e 2 La (f ◦ gϵ ) (s) = Jϵ (s)b(s, a) + cϵ (s, a) · ∇f (xϵ ) + tr Q (s, a) ∇ f (x ) ϵ ϵ . 2 2 On the other hand, the right-hand side of (57) is 1 Lhϵ (a) f (xϵ ) = b(xϵ , hϵ (a)) · ∇f (xϵ ) + tr Q(xϵ , hϵ (a)) ∇2 f (xϵ ) . 2 24

(58)

(59)

Since (57) holds for all f ∈ C 2 (S), comparing (58) and (59) yields uϵ · ∇f (xϵ ) + where

1 tr(Mϵ ∇2 f (xϵ )) = 0 2

for all f ∈ C 2 (S),

1 uϵ := Jϵ (s)b(s, a) + cϵ (s, a) − b(xϵ , hϵ (a)), 2

and

Mϵ := Jϵ (s)Q(s, a)Jϵ (s)⊤ − Q(xϵ , hϵ (a)).

Since the identity holds for all f ∈ C 2 (S), take local quadratic test functions f (x) = p · (x − xϵ ) + 21 (x − xϵ )⊤ H(x − xϵ ), multiplied by a smooth cutoff if needed. Thus ∇f (xϵ ) = p and ∇2 f (xϵ ) = H can be chosen arbitrarily. Choosing H = 0 and arbitrary p gives uϵ = 0; choosing p = 0 and arbitrary symmetric H gives Mϵ = 0. uϵ = 0 and Mϵ = 0. Therefore, for each ϵ, Q(gϵ (s), hϵ (a)) = Jϵ (s)Q(s, a)Jϵ (s)⊤ , and

1 b(gϵ (s), hϵ (a)) = Jϵ (s)b(s, a) + cϵ (s, a), 2

(cϵ )k = tr Q(s, a)∇2s gϵk (s) .

(60) (61)

We now differentiate (60) and (61) at ϵ = 0. Smoothness of the flows gives gϵ (s) = s + ϵX(s) + o(ϵ),

hϵ (a) = a + ϵY (a) + o(ϵ),

Jϵ (s) = I + ϵ∇s X(s) + o(ϵ),

∇2s gϵk (s) = ϵ∇2s Xk (s) + o(ϵ).

For the diffusion coefficient, Q(gϵ (s), hϵ (a)) = Q(s, a) + ϵ ∇s Q(s, a)[X(s)] + ∇a Q(s, a)[Y (a)] + o(ϵ), whereas Jϵ (s)Q(s, a)Jϵ (s)⊤ = Q(s, a) + ϵ (∇s X(s))Q(s, a) + Q(s, a)(∇s X(s))⊤ + o(ϵ). Equating the O(ϵ) terms gives ∇s Q(s, a)[X(s)] − (∇s X(s))Q(s, a) − Q(s, a)(∇s X(s))⊤ + ∇a Q(s, a)[Y (a)] = 0, which is (14). For the drift coefficient, b(gϵ (s), hϵ (a)) = b(s, a) + ϵ ∇s b(s, a)X(s) + ∇a b(s, a)Y (a) + o(ϵ), whereas 1 1 Jϵ (s)b(s, a) + cϵ (s, a) = b(s, a) + ϵ (∇s X(s))b(s, a) + ∆Q(s,a) X(s) + o(ϵ). 2 2 Equating the O(ϵ) terms gives 1 ∇s b(s, a)X(s) − (∇s X(s))b(s, a) + ∇a b(s, a)Y (a) − ∆Q(s,a) X(s) = 0, 2 which is (13). This completes the proof. F.4

Proof of Theorem 3.4

Proof. Fix (s, a) ∈ Ω and fix |ϵ| < ϵ0 small enough so that all points that appear below stay in Ω (which is possible by openness of Ω and continuity of (ϵ, s, a) 7→ (gϵ (s), hϵ (a))). We first prove reward invariance, then generator equivariance. Reward invariance follows by differentiating along the flow. Define ρ(ϵ) := r(gϵ (s), hϵ (a)). Using d d dϵ gϵ (s) = X(gϵ (s)) and dϵ hϵ (a) = Y (hϵ (a)), the chain rule gives ρ′ (ϵ) = ∇s r(gϵ (s), hϵ (a)) · X(gϵ (s)) + ∇a r(gϵ (s), hϵ (a)) · Y (hϵ (a)). 25

Since (12) holds on Ω, the right-hand side is 0 whenever (gϵ (s), hϵ (a)) ∈ Ω. Hence ρ′ (ϵ) = 0 for all such ϵ, so ρ(ϵ) ≡ ρ(0) = r(s, a) and thus r(gϵ (s), hϵ (a)) = r(s, a), i.e. (9) holds on Ω. For generator equivariance, define sϵ := gϵ (s),

aϵ := hϵ (a),

Jϵ (s) := ∇s gϵ (s).

We first derive finite coefficient transport identities. Since (14) holds along (sϵ , aϵ ), the matrix Q̄ϵ := Q(gϵ (s), hϵ (a)) satisfies

Q̄′ϵ = (∇s X(sϵ ))Q̄ϵ + Q̄ϵ (∇s X(sϵ ))⊤ ,

Q̄0 = Q(s, a).

The matrix

Q̂ϵ := Jϵ (s)Q(s, a)Jϵ (s)⊤ satisfies the same ODE and the same initial condition. Hence Q(gϵ (s), hϵ (a)) = Jϵ (s)Q(s, a)Jϵ (s)⊤ .

Similarly, using (13) and the Hessian evolution of gϵ , one obtains 1 b(gϵ (s), hϵ (a)) = Jϵ (s)b(s, a) + cϵ (s, a), 2

(cϵ )k = tr Q(s, a)∇2s gϵk (s) .

Now let f ∈ C 2 (S). By the chain rule, ∇(f ◦ gϵ )(s) = Jϵ (s)⊤ ∇f (gϵ (s)), and ∇2 (f ◦ gϵ )(s) = Jϵ (s)⊤ ∇2 f (gϵ (s))Jϵ (s) +

d X

∂k f (gϵ (s)) ∇2s gϵk (s).

k=1

Therefore, (La (f ◦ gϵ ))(s) = Jϵ b(s, a) + 21 cϵ (s, a) · ∇f (gϵ (s)) 1 + tr Jϵ Q(s, a)Jϵ⊤ ∇2 f (gϵ (s)) 2 1 = b(gϵ (s), hϵ (a)) · ∇f (gϵ (s)) + tr Q(gϵ (s), hϵ (a))∇2 f (gϵ (s)) 2 = (Lhϵ (a) f )(gϵ (s)). Thus

La Ugϵ = Ugϵ Lhϵ (a) . Together with reward invariance already proved, (gϵ , hϵ ) is an exact value-preserving structure on Ω. F.5

Proof of Theorem 3.5

Proof. Fix a point s0 in the neighborhood where the hypotheses hold. Because D is a smooth distribution of constant rank K, there exists a (possibly smaller) neighborhood U of s0 and smooth vector fields X1 , . . . , XK on U such that D(s) = span{X1 (s), . . . , XK (s)} for all s ∈ U and {X1 (s), . . . , XK (s)} is linearly independent for every s ∈ U . Indeed, in any local chart, D is represented by a rank-K smooth subbundle of T U , and constant rank allows choosing a smooth local frame. Since the Xi span D and D is involutive, for each i, j there exist smooth coefficient functions c1ij , . . . , cK ij on U such that K X [Xi , Xj ] = cℓij Xℓ on U. (62) ℓ=1

Because X1 (s0 ) ̸= 0, after shrinking U if necessary we may assume X1 (s) ̸= 0 for all s ∈ U . We will straighten X1 and then reduce the rank by one on a transversal, using involutivity to guarantee that the reduction is well-defined. 26

We invoke the flow-box (straightening) theorem for a nonvanishing smooth vector field. There exist local coordinates ϕ : U → V ⊂ Rd , ϕ(s) = (y 1 , . . . , y K , z 1 , . . . , z d−K ), with ϕ(s0 ) = 0, such that in these coordinates the pushforward of X1 is the constant coordinate field: ϕ∗ X1 =

∂ ∂y 1

on V.

(63)

Equivalently, writing s = ϕ−1 (y, z) and abusing notation by identifying vector fields with their coordinate expressions, (63) means X1 = ∂y1 in (y, z)-coordinates. Write the remaining generators X2 , . . . , XK in these coordinates as Xj =

d X

aα j (y, z)

α=1

∂ , ∂uα

u = (u1 , . . . , ud ) = (y 1 , . . . , y K , z 1 , . . . , z d−K ),

with smooth coefficients aα j . Since X1 = ∂y 1 , the Lie bracket with Xj becomes # " d d X ∂aα ∂ ∂ ∂ X α j , a (y, z) = (y, z) α , [X1 , Xj ] = j 1 α 1 ∂y α=1 ∂u ∂y ∂u α=1 because the coordinate vector fields commute and ∂y1 differentiates the coefficient functions. On the other hand, by involutivity and (62) we have [X1 , Xj ] =

K X

cℓ1j (y, z) Xℓ = c11j (y, z)

ℓ=1

X ∂ + cℓ1j (y, z) Xℓ . 1 ∂y ℓ=2

Comparing these two expressions is useful when restricting to the transversal hypersurface Σ := {(y, z) ∈ V : y 1 = 0}. At points of Σ, the component of [X1 , Xj ] tangent to Σ is determined by the components of X2 , . . . , XK tangent to Σ. This is the mechanism that makes the reduced distribution on Σ involutive. Define the projected (tangential) vector fields on Σ by restricting each Xj (j ≥ 2) to Σ and discarding the ∂y1 component: X̄j := Xj Σ − Xj y 1

∂ Σ ∂y 1 Σ

j = 2, . . . , K.

Then each X̄j is a smooth vector field on Σ taking values in T Σ. Let D̄ be the distribution on Σ spanned by X̄2 , . . . , X̄K ; it has rank K − 1 near 0 because X1 , . . . , XK are linearly independent and X1 is transverse to Σ. We claim that D̄ is involutive on Σ. To see this, take i, j ∈ {2, . . . , K}. Because Xi and Xj are tangent to D and D is involutive, [Xi , Xj ] is a linear combination of X1 , . . . , XK . Restricting to Σ, the tangential component of [Xi , Xj ] along Σ is therefore a linear combination of the tangential components of X2 , . . . , XK along Σ, i.e. a linear combination of X̄2 , . . . , X̄K . More explicitly, write on U K X [Xi , Xj ] = cℓij Xℓ ℓ=1

and restrict to Σ; since X1 = ∂y1 is normal to Σ and X̄ℓ is the tangential part of Xℓ for ℓ ≥ 2, PK the tangential projection of [Xi , Xj ] Σ equals ℓ=2 (cℓij |Σ ) X̄ℓ . But the tangential projection of [Xi , Xj ] Σ is exactly [X̄i , X̄j ] because brackets of tangent vector fields on Σ computed as restrictions agree with intrinsic brackets on Σ. Hence [X̄i , X̄j ] ∈ D̄ on Σ. Therefore D̄ is involutive and has constant rank K − 1 near 0 ∈ Σ. At this point we apply induction on the rank. For K = 1, the straightening theorem already produces coordinates (y 1 , z 1 , . . . , z d−1 ) with X1 = ∂y1 , and then each z α is constant along the integral curves of X1 because ∂y1 z α = 0. Assume the claim holds for rank K − 1 distributions in dimension d − 1. 27

Applying the induction hypothesis to the involutive constant-rank distribution D̄ on the (d − 1)dimensional manifold Σ, we obtain (d − 1) − (K − 1) = d − K smooth functions I¯1 , . . . , I¯d−K on a neighborhood Σ0 ⊂ Σ of 0 such that their differentials are linearly independent on Σ0 and X̄j I¯α = 0 on Σ0 , ∀j = 2, . . . , K, ∀α = 1, . . . , d − K. (64) Equivalently, I¯ = (I¯1 , . . . , I¯d−K ) is constant on the integral leaves of D̄ in Σ0 . We now extend these invariants off Σ along the flow of X1 = ∂y1 . Define, for (y, z) near 0, Iα (y, z) := I¯α (0, y 2 , . . . , y K , z 1 , . . . , z d−K ), α = 1, . . . , d − K, 1 i.e. Iα is independent of y and coincides with I¯α on Σ. By construction, X1 Iα =

∂Iα = 0. ∂y 1

(65)

(66)

We next show Xj Iα = 0 for j = 2, . . . , K. Fix j ∈ {2, . . . , K} and α. On Σ we have Xj = (Xj y 1 ) ∂y1 + X̄j by definition of X̄j , so using (65) (no y 1 -dependence) and (64), (Xj Iα ) Σ = (Xj y 1 ) Σ · (∂y1 Iα ) Σ + (X̄j I¯α ) = 0 + 0 = 0. Thus Xj Iα vanishes on Σ. To propagate this off Σ, differentiate Xj Iα along X1 = ∂y1 and use involutivity: X1 (Xj Iα ) = Xj (X1 Iα ) + [X1 , Xj ]Iα . The first term is zero by (66). For the bracket term, use (62): [X1 , Xj ]Iα =

K X

cℓ1j Xℓ Iα .

ℓ=1

We already have X1 Iα = 0 and we are studying Xj Iα for j ≥ 2. Consider the vector w(ϵ) ∈ RK−1 with components wj (ϵ) := (Xj Iα )(gϵ (s0 )) for j = 2, . . . , K, along the X1 -flow; the identity above implies w satisfies a linear ODE of the form w′ (ϵ) = A(ϵ) w(ϵ) with continuous coefficients determined by cℓ1j , and the initial condition at ϵ = 0 (i.e. on Σ) is w(0) = 0. Uniqueness of solutions to linear ODEs yields w(ϵ) ≡ 0 for small ϵ, hence Xj Iα = 0 in a neighborhood of 0 in V for all j = 2, . . . , K and all α. Together with (66), we have shown Xi Iα = 0

for all i = 1, . . . , K, α = 1, . . . , d − K

(67)

on a neighborhood U0 ⊂ U of s0 . The relations (67) mean precisely that each Iα is constant along any curve tangent to D, hence constant on each integral leaf of D in U0 . It remains to verify functional independence. In the coordinates (y, z) constructed above, the functions Iα depend only on (y 2 , . . . , y K , z) through (65), ¯ Since extending by keeping them and their differentials are independent on Σ0 by construction of I. independent of y 1 does not create new linear dependencies among the differentials, the covectors dI1 , . . . , dId−K remain linearly independent on a neighborhood of 0 in V . Equivalently, the map I = (I1 , . . . , Id−K ) has rank d − K and thus provides (d − K) functionally independent invariants, constant on the integral leaves. Writing z := I(s) gives local intrinsic coordinates for the orbit/leaf space, without the need to explicitly form a quotient S/G. F.6

Proof of Theorem 3.6

Lemma F.3 (Transport of the HJB under an exact value-preserving structure). Assume Assumption B.9. Fix any ϑ ∈ G and let (gϑ , hϑ ) be an exact value-preserving structure in the sense of Definition 3.1. Suppose u is a (bounded) solution of the stationary HJB equation in Assumption B.9, i.e., n o β u(s) = max r(s, a) + (La u)(s) , ∀s ∈ S, (68) a∈A

in the solution concept stipulated by Assumption B.9 (e.g., classical C 2 or viscosity). Define the transformed function uϑ : S → R by uϑ (s) := u(gϑ (s)).

(69)

Then uϑ is also a solution of (68). In particular, if Assumption B.9 guarantees uniqueness of the (bounded) solution to (68), then uϑ ≡ u and hence u(gϑ (s)) = u(s) for all s. 28

Proof. Fix ϑ ∈ G and define uϑ by (69). We show that uϑ satisfies (68) pointwise. Let s ∈ S be arbitrary and write s′ := gϑ (s). Since u solves (68), at the point s′ we have n o ′ ′ a′ ′ β u(s′ ) = max r(s , a ) + (L u)(s ) . (70) ′ a ∈A

Now use that hϑ : A → A is a bijection (as required by the symmetry definition), so the set {hϑ (a) : a ∈ A} equals A. Therefore, the maximization over a′ ∈ A can be re-parameterized by a′ = hϑ (a), yielding the exact identity n o n o ′ ′ a′ ′ ′ hϑ (a) ′ max r(s , a ) + (L u)(s ) = max r(s , h (a)) + (L u)(s ) . (71) ϑ ′ a ∈A

a∈A

′

Substituting (71) into (70) and recalling s = gϑ (s) gives n o β u(gϑ (s)) = max r(gϑ (s), hϑ (a)) + (Lhϑ (a) u)(gϑ (s)) . a∈A

(72)

Next we invoke the two exact-symmetry equalities from Definition 3.1. First, reward invariance gives, for every a ∈ A, r(gϑ (s), hϑ (a)) = r(s, a). (73) Second, generator equivariance gives, for every sufficiently smooth test function f (in particular for f = u in the classical case, or for the test functions used in the viscosity definition), (Lhϑ (a) f ) gϑ (s) = La (f ◦ gϑ ) (s). (74) Applying (8) with f = u and using (69) (so that u ◦ gϑ = uϑ ) yields, for every a ∈ A, (Lhϑ (a) u) gϑ (s) = La (u ◦ gϑ ) (s) = La uϑ (s).

(75)

Substituting (9) and (75) into (72) gives n o β u(gϑ (s)) = max r(s, a) + (La uϑ )(s) . a∈A

(76)

Finally, by definition (69), the left-hand side equals β uϑ (s). Hence (76) is exactly n o β uϑ (s) = max r(s, a) + (La uϑ )(s) , ∀s ∈ S, a∈A

which is (68) for uϑ . If Assumption B.9 asserts uniqueness of the (bounded) HJB solution, we conclude uϑ ≡ u, i.e., u(gϑ (s)) = u(s) for all s. Proof of Theorem 3.6. Fix ϑ ∈ G and define Vϑ (s) := V ⋆ (gϑ (s)). By Assumption B.9, the optimal value V ⋆ is a (bounded) solution of the HJB equation (68) (with u = V ⋆ ) and, moreover, is the unique (bounded) solution in the solution class specified by Assumption B.9. Since (gϑ , hϑ ) is an exact value-preserving structure, Lemma F.3 applies with u = V ⋆ and yields that Vϑ is also a solution of the same HJB equation. By uniqueness, we must have Vϑ ≡ V ⋆ , namely V ⋆ (gϑ (s)) = V ⋆ (s),

∀s ∈ S.

Because ϑ ∈ G was arbitrary, this establishes (18). For the invariant-coordinate representation, work on a neighborhood U ⊂ S where Theorem 3.5 provides a complete invariant coordinate map I : U → Rk in the stated sense: if s1 , s2 ∈ U and I(s1 ) = I(s2 ), then there exists ϑ ∈ G such that s2 = gϑ (s1 ). Define a function V̄ ⋆ on the set I(U ) by the rule V̄ ⋆ (z) := V ⋆ (s)

for any s ∈ U with I(s) = z.

(77)

We first check that (77) is well-defined. Indeed, if s1 , s2 ∈ U both satisfy I(s1 ) = I(s2 ) = z, completeness of I implies s2 = gϑ (s1 ) for some ϑ ∈ G, and the already-proved invariance (18) implies V ⋆ (s2 ) = V ⋆ (gϑ (s1 )) = V ⋆ (s1 ), so the value assigned by (77) does not depend on the choice of representative s. Finally, for any s ∈ U , taking z = I(s) and choosing s itself as the representative in (77) gives V̄ ⋆ (I(s)) = V ⋆ (s), i.e., V ⋆ (s) = V̄ ⋆ (I(s)) for all s ∈ U, which is the desired local factorization through invariant coordinates. 29

F.7

Proof of Corollary 3.7

Lemma F.4 (Hamiltonian invariance at the HJB maximand). Fix ϑ ∈ G. Assume reward invariance (9) and that the generator/action intertwining is in the (evolution-then-transform equals transformthen-evolution) pointwise form: for all a ∈ A, all f ∈ C 2 (S), and all s ∈ S, (Lhϑ (a) f ) gϑ (s) = La (f ◦ gϑ ) (s). (78) Then for the optimal value V ⋆ (which is invariant by Theorem 3.6), the HJB maximand H(s, a) := r(s, a) + (La V ⋆ )(s) satisfies, for all (s, a), H gϑ (s), hϑ (a) = H(s, a).

(79)

Proof. Fix (s, a) and ϑ. By reward invariance (9), r gϑ (s), hϑ (a) = r(s, a). For the generator term, apply (78) with f = V ⋆ : (Lhϑ (a) V ⋆ ) gϑ (s) = La (V ⋆ ◦ gϑ ) (s). By Theorem 3.6, V ⋆ ◦ gϑ = V ⋆ pointwise on S, hence La (V ⋆ ◦ gϑ ) (s) = La V ⋆ (s). Adding the reward identity yields (79). Proof of Corollary 3.7. Fix ϑ ∈ G and s ∈ S. Define H(s, a) := r(s, a) + (La V ⋆ )(s), so that Argmax(s) = arg max H(s, a). a∈A

Let a ∈ Argmax(s). We must show hϑ (a) ∈ Argmax(gϑ (s)), i.e., H gϑ (s), hϑ (a) ≥ H gϑ (s), a′ ∀a′ ∈ A. Take an arbitrary a′ ∈ A. Because hϑ : A → A is a bijection, there exists a unique ã ∈ A such that a′ = hϑ (ã). By Lemma F.4, H gϑ (s), a′ = H gϑ (s), hϑ (ã) = H(s, ã), H gϑ (s), hϑ (a) = H(s, a). Since a ∈ Argmax(s), we have H(s, a) ≥ H(s, ã) for every ã ∈ A, hence in particular H gϑ (s), hϑ (a) = H(s, a) ≥ H(s, ã) = H gϑ (s), a′ . Because a′ was arbitrary, this proves hϑ (a) ∈ Argmax(gϑ (s)). For the final sentence, note that the implication just proved shows the set-valued map Argmax(·) is orbit-consistent: if a⋆ (s) ∈ Argmax(s) then hϑ (a⋆ (s)) ∈ Argmax(gϑ (s)). Therefore, on any neighborhood where one can choose a selector that is additionally equivariant (i.e., satisfies a⋆ (gϑ (s)) = hϑ (a⋆ (s))), that selector is automatically optimal pointwise because it selects from Argmax(·) everywhere on that neighborhood. F.8

Proof of Theorem 4.2

Proof. Fix ϑ ∈ G and write g := gϑ , h := hϑ . Let W (s) := (V ⋆ ◦ g)(s) = V ⋆ (g(s)). Define the HJB maximand (Hamiltonian) operator n o H[U ](s) := sup r(s, a) + (La U )(s) . a∈A ⋆

By Assumption B.9 and Theorem 2.1, V is a bounded (classical) solution of βV ⋆ (s) = H[V ⋆ ](s), 30

∀s ∈ S.

(80)

We first show that W satisfies an approximate HJB equation with a uniform residual. Set δ := εr + εL ∥V ⋆ ∥C 2 . Take an arbitrary state s ∈ S. Evaluate the HJB identity (80) at the transformed state g(s): n o ′ βV ⋆ (g(s)) = sup r(g(s), a′ ) + (La V ⋆ )(g(s)) .

(81)

a′ ∈A

Because h : A → A is a bijection (group action on actions), we may reparameterize the supremum by a′ = h(a): n o n o ′ sup r(g(s), a′ ) + (La V ⋆ )(g(s)) = sup r(g(s), h(a)) + (Lh(a) V ⋆ )(g(s)) . (82) a′ ∈A

a∈A ⋆

Substituting (82) into (81) and using W (s) = V (g(s)) yields n o βW (s) = sup r(g(s), h(a)) + (Lh(a) V ⋆ )(g(s)) .

(83)

a∈A

Now apply approximate reward invariance (20): for every a ∈ A, −εr ≤ r(g(s), h(a)) − r(s, a) ≤ εr .

(84)

Next apply the approximate generator equivariance at the test function f = V ⋆ (recall V ⋆ ∈ C 2 (S) ∩ F). In the pointwise pushed-forward form consistent with the HJB transport used earlier, for every a ∈ A and s ∈ S we have (Lh(a) V ⋆ )(g(s)) − (La (V ⋆ ◦ g))(s) ≤ εL ∥V ⋆ ∥C 2 = δ − εr .

(85)

Combining (84) and (85) gives, for every a, r(g(s), h(a)) + (Lh(a) V ⋆ )(g(s)) ≤ r(s, a) + (La W )(s) + δ, and also the reverse inequality r(g(s), h(a)) + (Lh(a) V ⋆ )(g(s)) ≥ r(s, a) + (La W )(s) − δ. Taking the supremum over a ∈ A in both displays and using (83) yields the two-sided residual bound H[W ](s) − δ ≤ βW (s) ≤ H[W ](s) + δ,

∀s ∈ S,

(86)

equivalently βW (s) − H[W ](s) ≤ δ,

∀s ∈ S.

We now convert the residual inequality (86) into the stated uniform bound on W − V ⋆ . The argument is a discounted maximum principle comparison; we include it as a lemma for completeness. Lemma F.5 (Residual comparison for discounted HJB under ellipticity). Assume b, Q, r satisfy Assumption B.9(i)–(ii) and let U, V ∈ C 2 (S) be bounded. Assume V solves βV = H[V ] pointwise and U satisfies |βU − H[U ]| ≤ δ pointwise. Then ∥U − V ∥∞ ≤

δ . β

Proof. Let M := sups∈S (U (s)−V (s)) and fix η > 0. Choose sη such that U (sη )−V (sη ) ≥ M −η. Define the penalized function Φη (s) := U (s) − V (s) − η ∥s − sη ∥2 . Under Assumption B.9(iii), we may work on a neighborhood where Φη attains a maximum at some s̄η . At s̄η we have the first- and second-order conditions ∇2 Φη (s̄η ) ⪯ 0,

∇Φη (s̄η ) = 0, hence ∇(U − V )(s̄η ) = 2η(s̄η − sη ), 31

∇2 (U − V )(s̄η ) ⪯ 2ηI.

Fix any a ∈ A and consider the linear operator La in Itô form La f = b(·, a)·∇f + 12 tr(Q(·, a)∇2 f ). Using boundedness of b and Q from Assumption B.9(i) and Q ⪰ 0, we obtain at s̄η (La (U − V ))(s̄η ) = b(s̄η , a) · ∇(U − V )(s̄η ) +

1 tr Q(s̄η , a)∇2 (U − V )(s̄η ) ≤ C1 η, 2

where C1 is a constant independent of η (coming from ∥b∥∞ and ∥Q∥∞ ), because ∥∇(U −V )(s̄η )∥ = 2η∥s̄η − sη ∥ and ∇2 (U − V )(s̄η ) ⪯ 2ηI imply the drift term is O(η) and the diffusion term is ≤ η tr(Q(s̄η , a)) ≤ Cη. Therefore, for every a, r(s̄η , a) + (La U )(s̄η ) = r(s̄η , a) + (La V )(s̄η ) + (La (U − V ))(s̄η ) ≤ r(s̄η , a) + (La V )(s̄η ) + C1 η. Taking the supremum over a yields H[U ](s̄η ) ≤ H[V ](s̄η ) + C1 η = βV (s̄η ) + C1 η. Using the residual upper bound βU ≤ H[U ] + δ at s̄η gives βU (s̄η ) ≤ H[U ](s̄η ) + δ ≤ βV (s̄η ) + δ + C1 η, hence U (s̄η ) − V (s̄η ) ≤

δ C1 + η. β β

Because s̄η maximizes Φη , we have M − η ≤ U (sη ) − V (sη ) ≤ U (s̄η ) − V (s̄η ), so

C1 δ + η. β β Letting η ↓ 0 yields M ≤ δ/β, i.e. sup(U − V ) ≤ δ/β. Applying the same argument to V − U (using |βU − H[U ]| ≤ δ) yields sup(V − U ) ≤ δ/β. Therefore ∥U − V ∥∞ ≤ δ/β. M −η ≤

Apply Lemma F.5 with U = W and V = V ⋆ . The residual bound (86) gives |βW − H[W ]| ≤ δ and (80) gives βV ⋆ = H[V ⋆ ], hence δ 1 ∥W − V ⋆ ∥∞ ≤ = εr + εL ∥V ⋆ ∥C 2 . β β Recalling W = V ⋆ ◦ gϑ proves (21). F.9

Proof of Theorem B.2

Proof. Let w := (θ, ϕ) ∈ Rm collect all parameters and write F (w) := Lsym (θ, ϕ). The SGD iterate is wt+1 = wt − η gt ,

E[gt | wt ] = ∇F (wt ),

E ∥gt − ∇F (wt )∥22 | wt ≤ σ 2 .

Since F is L-smooth, for any x, y we have the standard descent inequality L ∥y − x∥22 . (87) 2 Apply (87) with x = wt and y = wt+1 = wt − ηgt , so that y − x = −ηgt . This gives the one-step bound L F (wt+1 ) ≤ F (wt ) + ⟨∇F (wt ), −ηgt ⟩ + ∥ − ηgt ∥22 2 Lη 2 = F (wt ) − η⟨∇F (wt ), gt ⟩ + ∥gt ∥22 . (88) 2 Take conditional expectation w.r.t. wt . By unbiasedness, E ⟨∇F (wt ), gt ⟩ | wt = ⟨∇F (wt ), E[gt | wt ]⟩ = ⟨∇F (wt ), ∇F (wt )⟩ = ∥∇F (wt )∥22 . (89) F (y) ≤ F (x) + ⟨∇F (x), y − x⟩ +

For the second moment, expand ∥gt ∥22 around ∇F (wt ): E ∥gt ∥22 | wt = E ∥(gt − ∇F (wt )) + ∇F (wt )∥22 | wt = E ∥gt − ∇F (wt )∥22 | wt + 2 E ⟨gt − ∇F (wt ), ∇F (wt )⟩ | wt + ∥∇F (wt )∥22 . (90) The cross term vanishes because E[gt − ∇F (wt ) | wt ] = 0, hence E ∥gt ∥22 | wt = ∥∇F (wt )∥22 + E ∥gt − ∇F (wt )∥22 | wt ≤ ∥∇F (wt )∥22 + σ 2 .

(91)

Taking conditional expectation of (88) and using (89)–(91) yields Lη 2 ∥∇F (wt )∥22 + σ 2 E[F (wt+1 ) | wt ] ≤ F (wt ) − η ∥∇F (wt )∥22 + 2 Lη 2 2 Lη ∥∇F (wt )∥22 + = F (wt ) − η 1 − σ . 2 2

(92)

1 Under η ≤ 1/L, we have 1 − Lη 2 ≥ 2 , so (92) implies

Lη 2 2 η σ . E[F (wt+1 ) | wt ] ≤ F (wt ) − ∥∇F (wt )∥22 + 2 2

(93)

Now take full expectation and rearrange: Lη 2 2 η E ∥∇F (wt )∥22 ≤ E F (wt ) − E F (wt+1 ) + σ . 2 2

(94)

Sum (94) from t = 0 to T − 1: −1 T −1 Lη 2 TX η X E ∥∇F (wt )∥22 ≤ E[F (wt )] − E[F (wt+1 )] + T σ2 2 t=0 2 t=0

= E[F (w0 )] − E[F (wT )] +

Lη 2 T σ2 . 2

(95)

Using the lower bound F (wT ) ≥ Linf almost surely, hence E[F (wT )] ≥ Linf , we obtain T −1 η X Lη 2 E ∥∇F (wt )∥22 ≤ F (w0 ) − Linf + T σ2 . 2 t=0 2

Divide by ηT /2 to get T −1 2(F (w0 ) − Linf ) 1 X E ∥∇F (wt )∥22 ≤ + ηLσ 2 , T t=0 ηT

which is exactly (28) (recall F = Lsym and w0 = (θ0 , ϕ0 )). For a rate statement, note that PT −1 min0≤t≤T −1 E∥∇F (wt )∥22 ≤ T1 t=0 E∥∇F (wt )∥22 . Choosing a stepsize on the order of η = √ Θ(1/ T ) (or using √ a standard diminishing schedule with comparable √ aggregate), the right-hand side becomes O(1/ T ), yielding min0≤t≤T −1 E∥∇F (wt )∥22 = O(1/ T ). F.10

Proof of Theorem B.4

b ω the learned Proof. Write (s, a) ∼ ρ and let ρS be the state marginal. For clarity, denote by rbω , bbω , Q models and define the (estimated) determining residuals (cf. (23)–(25)) brω,θ,ϕ (s, a) := Xθ (s) · ∇s rbω (s, a) + Yϕ (a) · ∇a rbω (s, a), R bω,θ,ϕ (s, a) := ∇sbbω (s, a) Xθ (s) − (∇s Xθ (s)) bbω (s, a) + ∇abbω (s, a) Yϕ (a) − 1 ∆ b R Xθ (s), b 2 Qω (s,a) bω,θ,ϕ (s, a) := ∇s Q b ω (s, a)[Xθ (s)]−(∇s Xθ (s)) Q b ω (s, a)−Q b ω (s, a) (∇s Xθ (s))⊤ +∇a Q b ω (s, a)[Yϕ (a)]. R Q

(N )

Let Lsym (θ, ϕ) denote the penalized objective in (26) when trained on a dataset of size N (with the corresponding fitted ω = ωN ), namely h i ) bωN ,θ,ϕ (s, a)∥22 + λQ ∥R bωN ,θ,ϕ (s, a)∥2F + λr |R brωN ,θ,ϕ (s, a)|2 L(N sym (θ, ϕ) =E(s,a)∼ρ ∥Rb Q (96) 2 + λnrm Es∼ρS ∥Xθ (s)∥22 − 1 . b ω⋆ = Q, rbω⋆ = r on supp(ρ), and By Assumption B.3, there exist (ω ⋆ , θ⋆ , ϕ⋆ ) such that bbω⋆ = b, Q (Xθ⋆ , Yϕ⋆ ) satisfies the determining equations on supp(ρ). Equivalently, on supp(ρ), ⋆

⋆

Rbθ ,ϕ (s, a) = 0,

Rrθ ,ϕ (s, a) = 0,

⋆

θ ,ϕ RQ (s, a) = 0,

b ω , rbω ). where R·θ,ϕ denotes the same expressions as above but with the true (b, Q, r) in place of (bbω , Q Also, by rescaling the generator (or by the normalization penalty), one can satisfy Es∼ρS ∥Xθ⋆ (s)∥22 = 1, hence ) inf lim sup L(N (97) sym (θ, ϕ) = 0. θ,ϕ

N →∞

Let (θN , ϕN ) be the parameters produced by the algorithm at dataset size N . The assumption that the symmetry objective can be optimized to its global minimum with vanishing optimization error means ) (N ) L(N sym (θN , ϕN ) − inf Lsym (θ, ϕ) −→ 0 θ,ϕ

as N → ∞.

(98)

Combining (97) and (98) yields ) L(N sym (θN , ϕN ) −→ 0.

(99)

(N )

Because every term in Lsym is nonnegative, (99) implies bωN ,θN ,ϕN (s, a)∥22 → 0, E(s,a)∼ρ ∥R bωN ,θN ,ϕN (s, a)∥2F → 0, E(s,a)∼ρ |R brωN ,θN ,ϕN (s, a)|2 → 0, E(s,a)∼ρ ∥R Q b (100) and also the normalization constraint is asymptotically satisfied: 2 Es∼ρS ∥XθN (s)∥22 − 1 → 0. (101) It remains to transfer (100) (estimated residuals) to the true determining residuals. This is the only place where we use “consistent dynamics estimation” in a mode strong enough to control the residuals, b ω , rbω and their first derivatives. which depend on bbω , Q Lemma F.6 (Consistency of residuals under consistent dynamics/derivative estimation). Assume that, on supp(ρ), the learned models satisfy rbωN → r, ∇s rbωN → ∇s r, ∇a rbωN → ∇a r, bbω → b, ∇sbbω → ∇s b, ∇abbω → ∇a b, N N N b ω → Q, ∇s Q b ω → ∇s Q, ∇a Q b ω → ∇a Q, Q N N N in ρ-mean square (or uniformly on supp(ρ)), and that the sequence of learned fields is uniformly bounded on supp(ρ) in the norms needed by the residuals, i.e. there exists C < ∞ such that for all N, h i E(s,a)∼ρ ∥XθN (s)∥22 + ∥∇s XθN (s)∥2F + ∥YϕN (a)∥22 ≤ C. Then brωN ,θN ,ϕN (s, a) − RrθN ,ϕN (s, a) E(s,a)∼ρ R 2

bωN ,θN ,ϕN (s, a)−RθN ,ϕN (s, a) → 0, E(s,a)∼ρ R b b 2

→ 0, 2

bωN ,θN ,ϕN (s, a)−RθN ,ϕN (s, a) → 0. E(s,a)∼ρ R Q Q F

Proof. We show the reward residual; the others follow by the same multilinear expansions. Write brωN ,θN ,ϕN (s, a)−RrθN ,ϕN (s, a) = Xθ (s)· ∇s rbω (s, a)−∇s r(s, a) +Yϕ (a)· ∇a rbω (s, a)−∇a r(s, a) . R N N N N 34

Using (u · v)2 ≤ ∥u∥22 ∥v∥22 and (x + y)2 ≤ 2x2 + 2y 2 , we obtain br − Rr 2 ≤ 2∥Xθ (s)∥22 ∇s rbω (s, a) − ∇s r(s, a) 2 + 2∥Yϕ (a)∥22 ∇a rbω (s, a) − ∇a r(s, a) 2 . R N N N N 2 2 Taking Eρ [·] and applying Cauchy–Schwarz to each term yields an upper bound by a product of (i) a uniform second-moment bound on ∥XθN ∥22 , ∥YϕN ∥22 and (ii) the mean-square estimation errors br − Rr |2 → 0. The of ∇s rbωN and ∇a rbωN , which vanish by the assumed consistency. Hence Eρ |R drift and diffusion residuals are finite sums of products of the same type (fields and their Jacobians times coefficient errors), so the same inequalities and consistency assumptions yield the stated convergences. Apply Lemma F.6 together with (100). For example, brωN ,θN ,ϕN (s, a)|2 + 2 Eρ R brωN ,θN ,ϕN (s, a) − RrθN ,ϕN (s, a) Eρ |RrθN ,ϕN (s, a)|2 ≤ 2 Eρ |R

−→ 0,

θN ,ϕN 2 and likewise Eρ ∥RbθN ,ϕN ∥22 → 0 and Eρ ∥RQ ∥F → 0. This is exactly the claimed “ρ-mean” convergence to the solution set of the determining equations on the data support: the sequence (XθN , YϕN ) drives the determining residuals to zero in ρ-mean square, hence satisfies (12)–(14) on supp(ρ) in the standard ρ-a.e. sense.

Finally, to connect to Definition 4.1, note that on supp(ρ), realizability and the preceding limit imply that the limiting infinitesimal fields solve the determining equations. By Theorem 3.4, exponentiating such a solution produces an exact local symmetry (gα , hα ) (for small |α|) on the same neighborhood, hence the corresponding reward and generator mismatches vanish there. Since the mismatches are continuous functionals of the coefficients and fields under the smoothness assumed throughout (and we have consistency of coefficient estimation plus vanishing determining residuals in ρ-mean), the induced mismatch levels on supp(ρ) satisfy εL → 0 and εr → 0 (in the natural “on-support” sense, e.g. ρ-essential supremum or ρ-mean square mismatch). F.11

Proof of Theorem B.6

Lemma F.7 (Discrete Grönwall for affine recursion). Let (en )n≥0 be nonnegative and satisfy en+1 ≤ (1 + γh) en + κhp+1 for some γ ≥ 0, κ ≥ 0, integer p ≥ 1, and stepsize h > 0. Then for every N ≥ 1, eN ≤ (1 + γh)N e0 + κhp+1

N −1 X

(1 + γh)j ≤ eγN h e0 +

j=0

with the convention e

γN h

−1

κ γN h e − 1 hp , γ

:= N h if γ = 0.

Proof. Iterating the inequality gives eN ≤ (1 + γh)N e0 + κhp+1

N −1 X

(1 + γh)j ,

j=0

which is the standard affine recursion bound. Since (1 + γh)N ≤ eγN h and the sum is a geometric series, N −1 X eγN h − 1 (1 + γh)N − 1 (1 + γh)j = ≤ , γh γh j=0 yielding the stated inequality. Proof of Theorem B.6. We prove the bound for the state flow; the action flow is identical with Xθ replaced by Yϕ and KS replaced by KA . Fix a bounded interval α ∈ [−A, A] and define the exact flow gα : S → S by the ODE d x(τ ) = Xθ (x(τ )), dτ

x(0) = s, 35

gα (s) := x(α).

By Assumption B.5, Xθ is Lipschitz on the relevant compact set containing the trajectories for |α| ≤ A, so existence and uniqueness of x(τ ) holds and gα (s) is well-defined for all s ∈ KS and |α| ≤ A. Let Φh denote the one-step numerical map of an order-p method with stepsize h applied to ẋ = Xθ (x). For each integer N ≥ 0, define the numerical iterate sn+1 = Φh (sn ),

s0 = s,

ĝnh (s) := sn ,

so that for a given α with |α| ≤ A we take N := ⌊|α|/h⌋ and (for simplicity of exposition) assume α = N h; the general case with a final partial step is treated by the same argument and only changes constants. Two standard ingredients are needed: a local truncation error bound and a stability (Lipschitz) bound for Φh on the compact region. Because the method has order p, there exists a constant κg (depending on p, bounds of derivatives of Xθ on the compact region, and method-specific coefficients) such that the one-step defect relative to the exact flow over time h satisfies sup x∈KS

Φh (x) − gh (x) 2 ≤ κg hp+1 ,

(102)

provided h is within the method’s admissible stability range on that region. Also, since Xθ is Lipschitz on KS and the method is a one-step scheme, there exists γg ≥ 0 such that ∥Φh (x) − Φh (y)∥2 ≤ (1 + γg h) ∥x − y∥2 ,

∀x, y ∈ KS ,

(103)

for all sufficiently small h in the solver stability region; here γg depends on the Lipschitz constant of Xθ on KS and on the method’s internal stability constants. Fix s ∈ KS and define the global error sequence at step n: en (s) := ∥sn − gnh (s)∥2 . We bound en+1 (s) in terms of en (s) using (102) and (103). Start from the identity sn+1 − g(n+1)h (s) = Φh (sn ) − gh (gnh (s)). Add and subtract Φh (gnh (s)) to split the difference: Φh (sn ) − gh (gnh (s)) = Φh (sn ) − Φh (gnh (s)) + Φh (gnh (s)) − gh (gnh (s)) . Taking norms and applying the triangle inequality gives en+1 (s) ≤ ∥Φh (sn ) − Φh (gnh (s))∥2 + ∥Φh (gnh (s)) − gh (gnh (s))∥2 . Because both the exact trajectory gnh (s) and the numerical iterate sn remain in the relevant compact region for 0 ≤ nh ≤ A (by Assumption B.5 and standard forward-invariance on compact sets for Lipschitz fields with small enough h), we may apply (103) to the first term and (102) to the second term: en+1 (s) ≤ (1 + γg h)en (s) + κg hp+1 . Since e0 (s) = ∥s0 − g0 (s)∥2 = 0, Lemma F.7 yields for N = α/h (with 0 ≤ α ≤ A) eN (s) ≤

κg γg α e − 1 hp . γg

Taking the supremum over s ∈ KS and over |α| ≤ A gives κ g sup sup ∥ĝα (s) − gα (s)∥2 ≤ eγg A − 1 hp =: Cg hp . γg |α|≤A s∈KS This establishes the first inequality in (29). The proof for ĥα is identical, with constants κh , γh depending on the Lipschitz constants of Yϕ on KA and the method stability characteristics, yielding supa∈KA ∥ĥα (a) − hα (a)∥2 ≤ Ch hp for |α| ≤ A. 36

F.12

Proof of Theorem B.7

(104)

Equivalently, if S ′ ∼ P (· | s, a) and S ′′ := g(S ′ ), then S ′′ ∼ P (· | g(s), h(a)). Proof. Fix (s, a) and define µ := P (· | s, a) and ν := P (· | g(s), h(a)). The assumption is exactly ν = g# µ, meaning ν(B) = µ(g −1 (B)) for all measurable B. By the defining property of the pushforward measure, for any bounded measurable f , Z Z f (x) ν(dx) = f (g(x)) µ(dx), S

which is precisely (104). Lemma F.9 (Bellman operator commutes with an exact value-preserving structure). Let the discounted optimal Bellman operator T be n o (T V )(s) := max r(s, a) + γ ES ′ ∼P (·|s,a) [V (S ′ )] , 0 < γ < 1. a∈A

Assume the exact value-preserving conditions hold for all (s, a): r(g(s), h(a)) = r(s, a),

P (· | g(s), h(a)) = g# P (· | s, a),

and that h : A → A is a bijection. Then for every bounded V and every s, (T V )(g(s)) = T (V ◦ g) (s).

(105)

Equivalently, writing (Ug V )(s) := V (g(s)), we have the operator identity Ug T = T Ug on bounded functions. Proof. Fix s ∈ S. Start from the definition of T at the transformed state g(s): n o ′ ′′ ′′ ∼P (·|g(s),a′ ) [V (S )] . (T V )(g(s)) = max r(g(s), a ) + γ E S ′ a ∈A

Because h is a bijection, every a′ ∈ A can be written uniquely as a′ = h(a) for some a ∈ A, hence n o (T V )(g(s)) = max r(g(s), h(a)) + γ ES ′′ ∼P (·|g(s),h(a)) [V (S ′′ )] . a∈A

Apply reward invariance r(g(s), h(a)) = r(s, a) to the first term. For the expectation term, apply Lemma F.8 with f = V to obtain ES ′′ ∼P (·|g(s),h(a)) [V (S ′′ )] = ES ′ ∼P (·|s,a) [V (g(S ′ ))] = ES ′ ∼P (·|s,a) [(V ◦ g)(S ′ )]. Substituting these two identities yields n o (T V )(g(s)) = max r(s, a) + γ ES ′ ∼P (·|s,a) [(V ◦ g)(S ′ )] = T (V ◦ g) (s), a∈A

which is (105). Lemma F.10 (Discounted optimal Bellman operator is a γ-contraction in ∥ · ∥∞ ). For any bounded V, W : S → R, ∥T V − T W ∥∞ ≤ γ ∥V − W ∥∞ . 37

Proof. Fix s ∈ S and define for each action a the quantities QV (s, a) := r(s, a) + γ E[V (S ′ ) | s, a],

QW (s, a) := r(s, a) + γ E[W (S ′ ) | s, a].

a∈A

By Jensen and boundedness, E (V − W )(S ′ ) | s, a ≤ E |(V − W )(S ′ )| | s, a ≤ ∥V − W ∥∞ . Hence |(T V )(s) − (T W )(s)| ≤ γ∥V − W ∥∞ for all s, and taking the supremum over s yields the claim. Proof of Theorem B.7. The statement that “the Bellman backup is invariant under augmentation by (g, h) in expectation” is a direct consequence of Lemma F.8. Concretely, suppose one draws a transition sample S ′ ∼ P (· | s, a) and forms the augmented sample Se′ := g(S ′ ) together with the augmented state-action pair (e s, e a) := (g(s), h(a)). For any bounded value function V , the one-step Bellman target at (e s, e a) is Ye := r(e s, e a) + γV (Se′ ) = r(g(s), h(a)) + γV (g(S ′ )). Taking conditional expectation given (s, a) and using reward invariance plus Lemma F.8 gives E[Ye | s, a] = r(s, a) + γ ES ′ ∼P (·|s,a) [V (g(S ′ ))] = r(g(s), h(a)) + γ ES ′′ ∼P (·|g(s),h(a)) [V (S ′′ )], so augmenting by (g, h) produces exactly the same expected backup as sampling directly from the true kernel at (g(s), h(a)). In this precise sense, exact value-preserving augmentation preserves the Bellman backup “in expectation”. Independently of augmentation, the (tabular) value-iteration recursion is Vk+1 = T Vk . By Lemma F.10, T is a γ-contraction in ∥·∥∞ , hence it has a unique fixed point V ⋆ satisfying V ⋆ = T V ⋆ . Moreover, ∥Vk − V ⋆ ∥∞ = ∥T k V0 − T k V ⋆ ∥∞ ≤ γ k ∥V0 − V ⋆ ∥∞ , which is the usual rate. For completeness, the exact value-preserving structure additionally implies an equivariance identity for the Bellman operator, given by Lemma F.9: (T V )(g(s)) = (T (V ◦ g))(s) for all V . This shows augmentation is not changing the underlying Bellman operator; it only replicates transitions consistently across the orbit. F.13

Proof of Theorem B.8

Lemma F.11 (Total variation controls expectation error). For probability measures µ, ν on S and any bounded measurable f : S → R, Z Z f d(µ − ν) ≤ ∥f ∥∞ dTV (µ, ν), dTV (µ, ν) := sup g d(µ − ν) . ∥g∥∞ ≤1

Proof. If ∥f ∥∞ = 0 the claim is trivial. Otherwise define g := f /∥f ∥∞ , so that ∥g∥∞ ≤ 1. Then Z Z Z f d(µ − ν) = ∥f ∥∞ g d(µ − ν) ≤ ∥f ∥∞ sup u d(µ − ν) = ∥f ∥∞ dTV (µ, ν). ∥u∥∞ ≤1

Lemma F.12 (Bellman max is 1-Lipschitz in the action-values). For any index set A and real families {xa }a∈A , {ya }a∈A , max xa − max ya ≤ max |xa − ya |. a∈A

a∈A

Proof. Let ax ∈ arg maxa xa . Then maxa xa = xax and maxa ya ≥ yax , hence max xa − max ya ≤ xax − yax ≤ max |xa − ya |. a

Swapping the roles of x and y yields maxa ya − maxa xa ≤ maxa |xa − ya |, and the claim follows. Lemma F.13 (T̃ is a γ-contraction in ∥ · ∥∞ ). Let Z Z n o n o (T V )(s) := max r(s, a)+γ V (s′ ) P (ds′ | s, a) , (T̃ V )(s) := max r̃(s, a)+γ V (s′ ) P̃ (ds′ | s, a) , a∈A

a∈A

with 0 < γ < 1 and r̃, P̃ any bounded reward and Markov kernel. Then for all bounded V, W , ∥T̃ V − T̃ W ∥∞ ≤ γ∥V − W ∥∞ . Proof. Fix s. For each a, define Q̃V (s, a) := r̃(s, a) + γ by Lemma F.12,

V dP̃ (· | s, a) and similarly Q̃W . Then

|(T̃ V )(s) − (T̃ W )(s)| = max Q̃V (s, a) − max Q̃W (s, a) a

≤ max |Q̃V (s, a) − Q̃W (s, a)| a Z = γ max (V − W ) dP̃ (· | s, a) .

(106)

s,a

which is exactly what (31) expresses when one sets, for each (s0 , a0 ), r̃ g(s0 ), h(a0 ) := r(s0 , a0 ), P̃ · | g(s0 ), h(a0 ) := g# P (· | s0 , a0 ). Indeed, |r̃(g(s0 ), h(a0 )) − r(g(s0 ), h(a0 ))| = |r(s0 , a0 ) − r(g(s0 ), h(a0 ))| ≤ εr and dTV (P̃ (· | g(s0 ), h(a0 )), P (· | g(s0 ), h(a0 ))) ≤ εP by (31). Since (g, h) is a bijection on state-action pairs, this implies the displayed sup-bounds over all (s, a). Fix any bounded V . For each (s, a), define the (one-step) action-values Z Z ′ ′ QV (s, a) := r(s, a) + γ V (s ) P (ds | s, a), Q̃V (s, a) := r̃(s, a) + γ V (s′ ) P̃ (ds′ | s, a). Then (T̃ V )(s) − (T V )(s) = max Q̃V (s, a) − max QV (s, a). a

Lemma F.12 gives |(T̃ V )(s) − (T V )(s)| ≤ max |Q̃V (s, a) − QV (s, a)|. a∈A

For a fixed (s, a), expand the difference: Z

V (s′ ) P̃ − P (ds′ | s, a)

V (s′ ) P̃ − P (ds′ | s, a) .

|Q̃V (s, a) − QV (s, a)| = r̃(s, a) − r(s, a) + γ ≤ |r̃(s, a) − r(s, a)| + γ

Substituting into the previous display gives, for every (s, a), |Q̃V (s, a) − QV (s, a)| ≤ εr + γ εP ∥V ∥∞ . Taking maxa and then sups yields the operator perturbation bound (32): ∥(T̃ V ) − (T V )∥∞ ≤ εr + γ εP ∥V ∥∞ . Lemma F.13 shows that T̃ is a γ-contraction in ∥·∥∞ , hence it admits a unique fixed point Ṽ ⋆ = T̃ Ṽ ⋆ . The true operator T is also a γ-contraction, with unique fixed point V ⋆ = T V ⋆ . To bound the fixed-point error, start from the identity Ṽ ⋆ − V ⋆ = T̃ Ṽ ⋆ − T V ⋆ . Add and subtract T̃ V ⋆ inside the right-hand side and apply the triangle inequality: ∥Ṽ ⋆ − V ⋆ ∥∞ ≤ ∥T̃ Ṽ ⋆ − T̃ V ⋆ ∥∞ + ∥T̃ V ⋆ − T V ⋆ ∥∞ . By Lemma F.13, ∥T̃ Ṽ ⋆ − T̃ V ⋆ ∥∞ ≤ γ∥Ṽ ⋆ − V ⋆ ∥∞ . Therefore, ∥Ṽ ⋆ − V ⋆ ∥∞ ≤ γ∥Ṽ ⋆ − V ⋆ ∥∞ + ∥T̃ V ⋆ − T V ⋆ ∥∞ , which rearranges to (1 − γ)∥Ṽ ⋆ − V ⋆ ∥∞ ≤ ∥T̃ V ⋆ − T V ⋆ ∥∞ . Applying (32) with V = V ⋆ gives ∥T̃ V ⋆ − T V ⋆ ∥∞ ≤ εr + γ εP ∥V ⋆ ∥∞ . Dividing by 1 − γ yields the claimed bound (33).

Record · ID 168362 · SHA-256 f0e9b8ae1168d34b
Retrieved via Conceptio — every document is proof-bundled with source, license, and retrieval metadata.