Geometry of Neural Reinforcement Learning in Continuous State and Action Spaces
Authors: Saket Tiwari
Omer Gottesman
George Konidaris
Venue: ICLR 2025 Oral (2025)
Keywords: reinforcement learning
deep learning
geometry
Primary area: reinforcement learning
TL;DR: We demonstrate, theoretically and empirically, that an RL agent with neural network policy induces a low-dimensional structure to the states sampled as its trajectories for deterministic environment

Abstract: Advances in reinforcement learning (RL) have led to its successful application in complex tasks with continuous state and action spaces. Despite these advances in practice, most theoretical work pertains to finite state and action spaces. We propose building a theoretical understanding of continuous state and action spaces by employing a geometric lens to understand the locally attained set of states. The set of all parametrised policies learnt through a semi-gradient based approach induce a set of attainable states in RL. We show that training dynamics of a two layer neural policy induce a low dimensional manifold of attainable states embedded in the high-dimensional nominal state space trained using an actor-critic algorithm. We prove that, under certain conditions, the dimensionality of this manifold is of the order of the dimensionality of the action space. This is the first result of its kind, linking the geometry of the state space to the dimensionality of the action space. We empirically corroborate this upper bound for four MuJoCo environments and also demonstrate the results in a toy environment with varying dimensionality. We also show the applicability of this theoretical result by introducing a local manifold learning layer to the policy and value function networks to improve the performance in control environments with very high degrees of freedom by changing one layer of the neural network to learn sparse representations.

Published as a conference paper at ICLR 2025
GEOMETRY OF NEURAL REINFORCEMENT LEARNING IN CONTINUOUS STATE AND ACTION SPACES
Saket Tiwari* Department of Computer Science Brown University
Omer Gottesman Amazon Web Services
George Konidaris
Department of Computer Science Brown University
ABSTRACT
Advances in reinforcement learning (RL) have led to its successful application in complex tasks with continuous state and action spaces. Despite these advances in practice, most theoretical work pertains to finite state and action spaces. We pro- pose building a theoretical understanding of continuous state and action spaces by employing geometric lens to understand the locally altained set of states. The set all parametrised policies learnt through a semi- aani-gradient based approach induces set of attainable states in RL. We show that the training dynamics of a two-layer neural policy induce a low dimensional manifold of attainable states embedded the high-dimensional nominal state space trained using an actor-critic algorithm. We prove that, under certain conditions, the dimensionality of this manifold is of the order of the dimensionality of the action space. This is the first result of its kind, linking the geometry of the state space to the dimensionality of the action space. We empirically corroborate this upper bound for four MuJoCo environ- ments and also demonstrate the results in a toy environment with varying dimen- sionality. We also show the applicability of this theoretical result by introducing a local manifold learning layer to the policy and value function networks to im- prove the performance in control environments with very high degrees of freedom by changing one layer of the neural network to learn sparse representations.
INTRODUCTION
The goal of reinforcement learning (RL) agent to learn policy that maximises its expected, time discounted cumulative reward (Sutton & Barto 1998). Recent advances in RL have lead to agents successfully learning in environments with enormous state spaces, such as games (Mnih) et al. 2015 Silver et al. (0166) Wurman al.] (022), robotic control in simulation (Lillicrap al., 2016 Schulman et al.) 2015] 2017) and real environments (Levine et al., 2016) Zhu et al. 2020 Deisenroth & Rasmussen 2011 Kaufmann et al. 2023). However, we do not have an understanding of the intrinsic complexity of these seemingly large problems.
We investigate the complexity of RL environments through geometric lens. We build the in- tuition behind the manifold hypothesis, which states that most figgh-dmeesiinall real-world datasets actually lie on close now-dimensionas manifolds (Tenenbaum 1997 Carlsson al. 2007 Fef- ferman et al. 2013; Bronstein et al., 2021); for example, the set of natural images is a very small, smoothly varying subset of all possible value assignments for the pixels. A promising geometric approach is to model this data as low-dimensional structure-a manifold-embedue dimensional ambient space. In supervised learning, especially deep learning theory, researchers have shown that the approximation error depends strongly on the dimensionality of the manifold (Sha-) ham et al. 2015 Pai et al.) 2019: Chen et al., 2019 Cloninger & Klock] 2020). thus connecting the learning complexity to the complexity of the underlying structure of the dataset. RL researchers have previously applied the manifold hypothesis - i.e. by hypothesizing that the effective state space lies low -ineeinnnio manifold (Smart & Kaelbling 22022] Mahadevan 20055: Machado ct al. 2017 2018 Banijamali et al.) 2018: Wu et al. 2019: Liu et al. 202 1), but the assumption has never been theoretically and empirically validated.
in a high-
"Corresponding author; saket tiwari @ brown.edu

---

Published as a conference paper at ICLR 2025
RL shares many similarities with control theory (Bertsekas 2012 framework, the objective is to drive the system, over time, to desired state or goal. Consequently, theoreticians and practitioners are often interested in the reachability of a control system to un- derstand what state is reachable given how system changes under control inputs, i.e. the system dynamics. Locally reachable states the set of states to which the system can possibly transition, starting from a fixed state, under all smooth time variant controls. Control theorists have long stud- ied the set of reachable states (Kalman 1960) using a differential geometric framework (Sussmann) 1973 1987). Theoretical research in the study of control systems is often focused on finding nec- essary and sufficient conditions in the system dynamies such that all states are reachable (Isidoni 1985 Sun et al. 2002, Respondek] 2005 Sun 2007) under all the admissible time-variant policies. In RL, the objective is to maximize the discounted return via gradient-based updates to the policy parameters, so the focus is on states attained through sequence of policies determined by these parameters.
2024
In a control-theoretic
Furthermore, theoretical understanding of states attainable using neural network (NN) policies gives us insight into the geometry and low-dimensional structure of data in RL. This requires utilis- ing an analytically tractable model of NNs. Ever since the remarkable success of neural networks, researchers have developed various theoretical models to better understand their efficacy. A theo- retical model intended to study complex object, such as neural network, often ends up making simplifying assumptions for tractability. One such theoretical model studies the evolution of neu- ral networks linearly in parameters during training of wide neural networks (Lee et al.) 2019 Jacot iet al. 2018), meaning in setting where the width approaches infinity. This has helped researchers develop theories of the generalization properties of neural networks (Jacot al.) 2018 Allen-Zhu et al. 2019a Wei et al., 2019 Adlam & Pennington 2020). We similarly utilize a single hidden layer neural network model for the policy that linear in terms of its parameters, not linear in the state, as the width approaches infinity, as has previously been applied to RL (Wang et al.) 2019 Cai
(et al. 2019a).
Within this theoretical framework, we provide a proof of the manifold hypothesis for deterministic continuous state and action RL environments with wide two-layer neural networks. We prove that the effective set of attainable states is subset of manifold and its dimensionality is upper bounded linearly in terms of the dimensionality of the action space, under appropriate assumptions, indepen- dent of the dimensionality of the nominal state space. The primary intuition is that the set of states locally attained are restricted by two factors: 1) the policy is time invariant and state dependent, and 2) the set of policies is constrained by the optimization of a wide, two-layer neural network using stochastic policy gradients. Our theoretical results are for deterministic environments with con- tinuous states and actions: we empirically corroborate the low-dimensional structure of attainable states on MuJoCo environments (Todorov et al.] (2012) by applying the dimensionality estimation algorithm by Facco et al. (2017). To show the applicability and relevance of our theoretical result, we empirically demonstrate that a policy can implicitly learn a low-dimensinn marginal computational overhead using the CRATE framework (Yu et al. 2023a by Pai et al.. 2024). We present algorithm that does two things simultaneously: 1) learns mapping to local low dimensional representation parameterised by a DNN, and 2) uses this effectively low-dimensional mapping to learn the policy and value function. Our modified neural network works out of the box with SAC (Haarnoja et al.] (0188) and we show significant improvements in high dimensional DM control environments (Tunyasuvunakool al. 2020)-
representation with
BACKGROUND AND MATHEMATICAL PRELIMINARIES
We first describe the continuous-time Markov decision process (MDP), which forms the foundation upon which our theoretical result is based. Then we provide mathematical background various ideas from the theory of manifolds that we employ in our proof.
2.1 CONTINUOUS- TIME REINFORCEMENT LEARNING
We first analyse continuous nattinuuus-time reinforcement learning in a deterministic Markoy decision process (MDP) defined by the tuple M (S, A, fr. Su, X) over time t E [). S C Rdd is the set of all possible states of the environment. A C dd is the rectangular set of actions available to the agent.
T : S x A x R S and T C is smooth function that determines the state transitions:
2

---

Published as a conference paper at ICLR 2025
a, T) is the state which the agent transitions when it takes the action a at state for the s' = =Tss, (ssss, time period т. Note that T(s, a, 0) = 8, which means that the agent's 's state remains unchanged if an action is applied for a duration of determined by the reward function fr f S -> R. 8, denotes the state the agent is at time and at is the action it takes at time f. So is the fixed initial state of the agent / = O, and the MDP terminates at & = T. The agent lacks access to f and fr. and can only observe states and rewards at given time t E [0, The agent's decision-making process is determined by its policy т S A. Simply put, the agent takes action πT (s) in state s. The agent's goal is to maximize the discounted return ((x) for e t (s))ddl where Setc T(s,, (s)), for small and for all e [O,T). We define the action tangent mapping, 9 S - Rdd,, for an MDP as
0. The reward obtained for reaching the state is (s),
T(s, a,),a,c) - 0T (8,
Val 0) = lim -
c
de
Intuitively, this captures the direction of change in the state s after taking an action a. We consider the family of control affine systems that represent wide range of control systems (Isidori 1985
Murray Hauser 1991; Tedrake 2023). such that & = g(s) + Σ ha(s)a, where & is the time derivative of the state, 9, h4 - Rd, are infinitely differentiable (or smooth) functions. Similarly, x xs)) = I [s)...... Tdg (s)] is the direction of change in the agent's state following
policy ı at state for an infinitesimally short time. The curve in the set of possible states, or the state-trajectory of the agent, is a differential equation whose integral form is:
2 = 80 + 1 g(67)) + Σ
(1)
This solution is also unique (Wiggins) 1989) for fixed start state, sp, and Lipschitz continuous policy, π. The above curve is smooth if the policy is also smooth. Therefore, given an MDP M and smooth deterministic policy π e II, the agent traverses a continuous time state-trajectory or curve MM,a [0, + S. The value function at time & for a policy π is the cumulative future reward starting at time
the u ((())== - fx(
Note that the objective function, [TT, is the same 1. (5o). Our specification is very similar to classical control and continuous time RI. (Cybenko) 1989: Doya (000a) but we define the tran- sitions, T, differently. More recently, researchers have developed the theory for continuous-time RL in a model-free setting with stochastic policies and dynamics (Wang et al. 2020 Jia & Zhou, 2022ap.
2.2 MANIFOLDS
In practice, MDPs have a low-dimensional underlying structure resulting in fewer degrees of free- dom than their nominal dimensionality. In the Chectah MujoCo environment, where the Cheetah is constrained to a plane, the goal of the RL agent is to learn policy to make the Cheetah move forward as fast as possible. The actions available to the agent are to provide torques at each of the 6 joints. For example, RL agent learning from control inputs for the Cheetah MuJoCo envi- ronment, one can "minimally" describe the cheetah's 's state by its "pose", velocity, and position as opposed the entirety of the input vector. The idea of a low-dimensional manifold embedded in high-dimensional state space formalises this.
A function h X Y, from one open subset X C R to another open subset Y C R is a diffeomorphism if f is bijective, and both h and h- are differentiable. Intuitively, a low dimensional surface embedded in high dimensional Euclidean space can be parameterised by differentiable mapping, and if this mapping is bijective we term it a diffeomorphism. Here, X is diffeomorphic to Y. A manifold is defined as follows (Guillemin & Pollack 1974; Boothby 1986] Robbin et al. 2011).
Definition 1. A subset M C R is called a smooth m dimensional submanifold of TR iff every point p E M has open neighborhood U R such that U n M is diffeomorphic subset O C Rm A diffeomorphism, @ : U M + O is called a coordinate chart of M and the inverse, v) = s-i: + U ( M is called a smooth parameterisation.
open
3

---

Published as a conference paper at ICLR 2025
We illustrate this with an example in Figure || Further note that a coordinate chart is called local to some point P E U C M the diffcomorphism property holds in the neighborhood U. It offers a local "flattening" of the local neighborhood. It is called global if it holds everywhere in M but not all manifolds have global chart (e.g.. Figure []. If M C R is a smooth non-empty m-manifold, then m k. reflecting the idea that manifold of lower equal dimension than its ambient space. A smooth curve 7y 1 + M is defined from an interval I C R to the manifold M as a function that infinitely differentiable for all The derivative of y at is denoted as (t). The of derivatives of the curve at time t. (t), for all possible smooth 7 forms a set that is called the tangent space (MM) at the point p. For a precise definition, see the appendix A Taking partial derivatives of tb with respect to each coordinate x², we obtain the vectors in R
Oeih Oa.
Ozd
Oal
These vectors span the tangent space My at the point p. Therefore, locally the manifold can be alternatively represented as the space spanned by the non-linear bases: Span( Oy
0u*
2.3 VECTOR FIELDS, LIB-SERIES, AND CONTROL THEORY
Curves and tangent spaces in manifolds natu- rally lead to vector fields. In the same way that curve represents how agent's state changes continuously, a vector field captures this change locally at every point of the state space. A tangent vector can be represented as X [u,,..., .., U. where each v, is a function.
ga
Definition 2. The vector field X is called a C vector field if. in any local coordinate chart (U.P) with coordinates (Z1...., Im), the com- ponents of X in the local basis dw: are C" functions. That in local coordinates, X can where Brg. each component function v. : U + is
Figure 1: The surface of open cylinder of unit radius, denoted by 52, in IR3 is a 2D man- ifold embedded
be written as X(x) = E U1 (z) 2.
a 3D space. More formally, S = {(x, y, 22 +y² 1, E (-h, h)} where the cylinder's 's height 2h. One can smoothly pa- rameterise S2 as (0, b) = (sin f, cos 0 b). The
We denote by (M) the set of all smooth vector fields on manifold M. The rate of change of function ʃ C (M) at point x along the vector field X is defined by
coordinate chart is ó(x, y, 2) = (sin x,2).
0f(x) dr.
Lx()=x(UU) = Σ 0
(3)
Associated with every such vector field X E (M) and zo E M is the integral curve: x(t). Intuitively, following along the direction X for time t. the curve starting from zo reaches the point (t).. The solution to the ODE with the starting condition z(0) = 2o is denoted as the exponential map e,x (zx). One imagine that vector fields have connection to policies in the way policy determines the direction of change. Therefore, it is an effective way to model the change in an agent's state given vector field and an arbitrary fixed starting state over a time period.
Taylor series help approximate complex functions with polynomials; analogously, we will use the Lie series of the exponential map. To define this expansion, we first recursively define the derivative of Lie T (f) LxxLX (Lxx (()) for l E N+ where Lx (.) is defined in equation
3 The Lie series of
the exponential map with f(x) = x is (Jurdjevic
1997 Cheng et al. 2011)
8
g
c (x) = + (X
+
1+ (x).
(4)
+ 1)!
4

---

Published as a conference paper at ICLR 2025
3 MODEL FOR LINEARISED WIDE TWO-LAYER NEURAL POLICY
An RL agent in the policy gradient framework (Sutton et al. 1999 Konda & Tsitsikiis 1999) is
equipped with policy и parameterised by parameters and takes gradient ascent steps the di- rection of / (T((()())) 0)). Suppose that the agent's 's policy is parameterised by a wide two-layer neural
network policy. This update direction can be estimated in different ways (Williams) 1992 Kakade 2001). Such an algorithm generates a sequence of parameters:
@ + 0 + (0),
(5)
where 1) the learning rate. In our setting, a neural RL agent parameterises the policy as two-layer neural network with smooth activation. We highlight the salient details below.
3.1 LINEAR PARAMETERISATION OF NEURAL POLICY
For the set of permissible policies we consider the family of two-layer feed forward neural networks with GeLU activation (Hendrycks & Gimpel 2016), which is smooth analog of the popular ReLU activation (Nair & Hinton| 2010). We follow the parameterisation for a two-layer fully connected neural network employed by Cai et al. (2019b) and Wang et al. (2019) for analysis of RL algorithms, which is also used in theoretical analyses of wide neural networks in supervised learning (Allen-Zhi) al., 20T9b Gao et al., [219 Lee al. 2019). For weight vector W of the first layer and weights C. for the last layer a shallow, two-layer, fully connected neural network is parameterised as:
f(:: J Σ CW 18)
Vn
where W Rad, is the vector of first layer parameters where each WR is a de length vec- tor block and therefore the complete vector W = 2 Wn), is GeLU activation, and B Id. xa matrix comprised of column vectors of dimension d, denoted Ck. meaning B B= [2, Cn). Here is the width of the neural network. The parameters are randomly initialised i.i.d as BR Unif(-1, 1) and W, Normal /d.dd)) where la. is a d, x d, iden- tity matrix and Unif is the uniform distribution. During training. Cai et al. (019b)) and Wang et al. (2019) only update W while keeping B fixed to its random initialisation despite which, for a slightly different policy gradient based learning, the agent learns near optimal policy. Researchers study neural networks in simplified theoretical settings to advance the understanding of a complex system while keeping the mathematics tractable (Li & Yuan 2017, Jacot et al.) 2018 Du et al. 2018 Mei et al, 2018a Allen--Zhu et al. 20195). While this shallow model of neural networks does away with complexity from multiple layers, captures the preererrametiizttioonnn in NNs.
Let 0 be the initial parameters of the policy network defined in Equation 6 of the policy defined as
A linear approximation
(s; W) = =s w + Vef(s;0)le= (W- - w - f(s; W0)
D(:; Wo) -WW) (7)
Chee (W s)sT] is a do x nd, fea- ture matrix for the input 8, sf is the gradient of GeLU function w.r.t the input, represents the dot product, and the matrix & formed by the concatenation of dg x de matrices Bup' s)sT for k = 1, I....,n.. n. This results in matrix of size da x ndg. W is an nd, vector described above. We will omit the parameters W, W and from the representation of policies when there is no ambiguity. It is a linear approximation because is linear in the weights W and non-linear, within T. in the initial weights W' and the state s. This leads us to the definition of the family of linearised policies for fixed initialisation W", similar to Wang et al. (2019).
where (2;; Wo) = - [CP''(WP . S)ST CD s)sT,
Definition 3. For a constant r > 0, and fixed Wo. For all widths E N20. we define
{
Fwo.r.n
Ch Cl s)s W. |W - W [r]
This linearised approximation of the policy simplifies our analysis of the set of reachable states. We further note that it might seem restrictive to consider a network without bias, but we can extend this analysis by adding another input dimension, which is always set to 1.
5

---

Published as a conference paper at ICLR 2025
3.2 CONTINUOUS TIME POLICY GRADIENT
Under the parameterisation described above, the sequence of neural net parameters as described by the updates in equation determined by the semi-gradient update direction
VoJ= =E V (s, a, tVo W)]
where the expectation is over the visitation measure Px and Q S A x [0, T] is the action-value function that represents the value of taking a constant action a at time l. For further details see Appendix H As is usually done, the following stochastic gradient based update rule approximates the true gradient for the policy parameters
B E
W(k+1)g Wk
7
N.. ab. Ln (s6: Wo).
where Way represents the parameters after gradient steps with learning rate y) Q Q is the action-
value function associated with policy parameterised by flin (: Whh), Bw., {88. ds. tb) b=1 is randomly chosen batch of data from samples of the SDE (Doya, 2000b, Jia & Zhou 2022a
dS. ((.)) + Σ (s) Wan)) di (St)dur,
We is the d2 dimensional Wiener process where a Rd. xd, is the exploration component of the agent. We assume access to an oracle that gives us the gradients T.QW.,, which do not need be true in practice. Therefore, sample Bw an i.i.d. set of samples from (I...., J" for large N', of size B, Thus we can write the expectation of the gradient update as follows
where Wo W (see Section 3.1
B Σ 6=11
N' N Σ
EBw
B
V
V.QWan Wan (ss, ab ( (8; Wo)
=
V
Wen (an ag (,,44) (s, Wo),
where we have an appropriate function Q such that the above condition is satisfied. Let the term on the right hand side be denoted by (WW in the limit N' + co. Here, is the exploration component of the dynamics. We re-write the update rule from equation as follows,
W (k+1)) Wky = 1) [WW[[WW)-ww + *n(WWAm kmy BW,..) = G(Wan) 7),
(9)
Whn (st, ab (o) (su;; Wo) ww (WW))WWw... w= There- fore, we have EBg [C(W, Bw)] )] O given an unbiased sampling mechanism for Bw Simiar formu- lation of SGD also used in supervised learning (Cheng al. 2020 Ben Arous et al. 2022).
where kWr Bw..,) 1/2 - V V.QWa.
MAIN RESULT: LOCALLY ATTAINABLE STATES
The state space is typically thought of as dense Euclidean space with all states reachable, but not necessarily the case that all such states are reachable by the agent. Three main factors constrain states available to an agent: 1) the transition function, 2) the family of functions to which policy belongs, and 3) the optimization process which determines the dynamics of parameters of the policies. We therefore are interested in the set of states attained by the trajectories of linearised policy with parameters that are optimised as in Section 3.2 around fixed state for time 8. The properties this set gives us a proxy for the "local manifold" around any arbitrary state.
A vector field, its exponential map, and the corresponding Lie series described in Section 2.3 analogous to parameterised policy, the state transition based on this policy, and an approximation of this rollout. To formalise this, denote the vector field determined by the parameters W of linearised policy with initialisation W is
X(W) g(x) ha, (a;; W²)W")o(r; W
(10)
The set of states attained by the rollout of this policy, parameterised by W, W over time 6 is X(W) (s), i.e. the image of the interval (0, under the exponential map corresponding therefore (0.5)

---

Published as a conference paper at ICLR 2025
to the vector field X (W). Moreover, We is randomly initialised, and the parameters W are obtained through stochastic semigradient updates (equation
There are two time scales: one is the time of policy rollouts and the other is the policy parameter optimisation. This complicates the analysis. We will use t for time in the physical sense of an RL environment and for the gradient update step. Continuous-ting
analogues for discrete stochas- tic gradient descent algorithms at small step sizes have yielded remarkable theoretical analyses of
algorithms (Mei et al.] 20188) Chizat & Bach] 2018 Jacot et al. 2018 Lee et al., 2019 Cheng et al., 20200 Ben Arous et al. 2022). Therefore, to analyse the evolution of the attainable states under dime-discretized sequence of parameters, we derive an approximate continuous time dynam- ies for the evolution of the randomly initialised parameters W. Many theoretical frameworks that study SGD continuous time seek to approximate the evolution of the high-dimensional parameter distribution, but we seek to closely approximate the Lie series. We therefore utilise the theoret- ical framework provided by [Ben Arous et al. (2022). with appropriate modifications, to analyse continuous-time dynamics of relevant statistics in the infinite-width limit.
Let En. G. the semi-gradients for linearised policy of width n, shn. Let 7/ be a sequence of learning rates such that Thn + 0 as - at rate _
Vn For random variable W, n. which
determines the distribution of the nd, parameters, et (WW))(s denote the purwaforward of W, n of the exponential map. In the case of random variables, the attained set of states is sampled from this time-dependent push-forward of the distribution WT, n. where T is the gradient time step. We make the following assumptions
Assumption 4. Suppose Bw) = (W, J.W.BB. Bw) - Gn (W) for any n and given compuct set (W)
K there exists constant OH,K such that Eew (WW, I n for W K, where L? is the 2-norm.
This assumption is relaxed version of the assumption on the variance of the gradient update (as- sumption 4.4) made by Wang et al. (2019). We make a further assumption about the Lipschitz continuity of H, and Gn. similar to Ben Arous et al. (2022).
Assumption 5. Gn is locally Lipschitz continuous in W.
Furthermore, we assume that the activation, sc has bounded first and second derivatives everywhere in R. This assumption holds for GeLU activation. We also denote by Jh, the d, x d, Jacobian of the d, x 1 vector-valued function hs))). We also define the proximity of a random variable to manifold in a probabilistic manner.
Definition 6. A random variable, X, is concentrated around a manifold M with a rate R. if Pr(distance(X, X, M)
0()) RR())
Intuitively, this means that the probability that the random variable X lies at more distance than D decays exponentially.
Theorem 1. Given a continuous time MDP M, a fixed state s, sequence of two-layer linearised neural network policy, flin, initialised with i.i.d samples from Normal(O, 1/d,), semi-gradiens based updates (Ma, En. Gn) which satisfy assumptions 4 5 then for varying öt E (0, 8) and fixed T > 0 the random variable defined by the push-jorward of the ran- dom variable WT, w.r. the exponential map Est (s) converges weakly a random variable S + 5 such that S concentrates around an m-dimensional manifold M& with m S 2dg 1 at rate R, that depends on the operator norms of the matrices Jhy(s), E (1...., ) the values 9x(s), k (L,...,del, т. Further, the variable S is O(83).
X(WW)
Intuitively, this means that in the infinite width limit for very low learning rates the probability mass of the furw-fo of the exponential map is concentrated around 2d, + dimensional manifold and this probability decays exponentially as one moves away from this manifold. The proof is provided in Appendix G The proof sketch is follows:
1. We expand the Lie series up to an error term of 83 (Appendix B
2. We then show the weak convergence of the dynamics of random variables that determine the Lie series in Appendix F this Section closely follows the proof by [en Arous et al.
(2022)
7

---

Published as a conference paper at ICLR 2025
3. Finally, show that the push forward of the random variable W through Lie series pansion is concentrated around a space spanned by 2d, + vectors for fixed t and therefore for variable there is 2d,, + around which the data lie, modulo the 6 distance.
Here, the directions in which the individual action dimensions change the state locally are h,,..., hd.. The mean second-order change: oj _k=1 Öse Σj²=1 ajJhy.. ( (ss), where Jhy is the Jacobian of the function (s) and az a constant that depends on the gradient time т, and the paraboloid. g first order partial derivative of g, and therefore tg + t²gg' is a paraboloid. This is similar to how local neighborhood is defined being spanned by bases vectors in Section 2.2 Informally extending and intuiting this result, one can hypothesize that over the training dynamics of a linearised neural network, if the parameters remain bounded, the union of trajectories starting from state 8 a "small" time interval then the trajectories are concentrated around a 2d + 3 manifold. The reason being that their is additional degrees of freedom from the gradient dynamics. This means "locally" the data is concentrated around some low-dimensional manifold whose dimensionality is linear in da-
The manifold M is locally derived being spanned by (hy, 3z 21,66 ty + t2g') locally
(h,(s)
5 EMPIRICAL VALIDATION
Our empirical validation is threefold. First, we show the validity of the linearised parameterisation of the policy (Equation 7 as a theoretical model for canonical NNs (Equation (6) Second, we verify that the bound on the manifold dimensionality in [heorem[hhldde in practice. In the third subsection, we demonstrate the practical relevance of our result by demonstrating the benefits of learning compact -iw-dimensional representations, without significant computational overhead.
5.1 APPROXIMATION ERROR WITH LINEARISED POLICY
We empirically observe the impact of our choice of lin- earised policies theoretical model for two-layer NNs by measuring the impact on the returns of this choice. We calculate the difference in returns for DDPG using canonical NNs and linearised NNs as parameterisations for its policy network, while only training the weights of the first layer. Let the empirically observed return to which the DDPG algorithm converges using a canonical NN policy be J and J!in be the same for a linearised policy. In figure 2 we report the value -J J(i on the y- axis and log, n on the x-axis for the Cheetah environment (Todorov et 2012: Brockman et al. (016). We present additional training curves in the appendix (figure) 7 that compare how the returns vary as training progresses. In- terestingly, at large widths (log2 > 15) the discounted returns match across training steps for canonical and linearised policy parameterisations. This suggests that the agent's 's learning dynamics captured by linearised policy
1250
Difference
750
500
250
10
11
12
13
15 16
17
(width)
Figure 2: We observe that the difference between the returns approaches zero as we increase the width.
oc. All results are averaged across 16 seeds.
5.2 EMPIRICAL DIMENSIONALITY ESTIMATION
We empirically corroborate our main result (Theorem in the MuJoCo domains provided in the OpenAl Gym (Brockman et al. (0166). These are all continuous state and action spaces with d, < d, for simulated robotic control tasks. The states are typically sensor measurements such as angles, velocities, or orientation, and the actions are torques provided at various joints. We estimate the dimensionality of the attainable set states upon training. To sample data from the manifold, we record the trajectories of multiple DDPG evaluation runs across different seeds (Lillicrap et al.) 2016). with two changes: we use GeLU activation (Hendrycks both policy and value networks, and we also use single hidden layer network instead of hidden layers for both networks. Performance is comparable to the original DDPG architecture (see Appendix D.. For background DDPG refer to Appendix [ These choices keep evaluation of
Gimpel 2016) instead of ReLū
8

---

Published as a conference paper at ICLR 2025
Freettied 2am
Predeteet Dm
Fredectdd
50020
0000
Nkm Samplen
Num Samplen
Num
Num Samples
(a) Walker2D
(b) Cheetah
(c) Reacher
(d) Swimmer
Figure 3: Estimated dimensionality of the attainable states, in blue, is far below de (green line) and also below 2d, + (red line) for four tasks, estimated using the method by Facco et al. (001)).
the upper bound as close to the theoretical assumptions possible while still resulting in reasonably good discounted returns. We then randomly sample states from the evaluation trajectories to obtain subsample of states, D {s, {s,);_1- We estimate the dimensionality with 10 different subsamples of the same size to provide confidence intervals.
We employ the dimensionality estimation algorithm introduced by Facco al.) (2017) 7017) which timates the intrinsic dimension of datasets characterized by non-uniform density and curvature, to empirically corroborate Theorem More details on the dimensionality estimation procedure are presented in the Appendix [ Estimates for four MuJoCo environments are shown in Figure 3 For all environments, the estimate remains below the limit of 2d, 1 in accordance with Theorem |
5.3 EMPIRICAL VALIDATION IN TOY LINEAR ENVIRONMENT
A deterministic system is fully reachable if given any start state, So E IRd,, the system can be driven to any goal state in Id. To contrast our result to classic control theory, we demonstrate that for a control environment which is fully reachable using time-variant the set of all the attainable states using a bounded family of linearised neural nets (definition is low-dimensional. common example of a fully reachable d, diddimensional linear control problem with ID controller is:
nstimated Intrincic Dimension
open loop policy
10
Extrinsic dimension
1
0
07
0
0
1
0
0
Figure 4: The intrinsic dimensionality estimate of attainable states linear fully reachable system under linearised pol- icy y-axis.
(())
s(1)
+
T ((),
(11)
which is fully reachable. This follows from the fact that a linear system à = Ax + Bu(t) is fully reachable if and only its controllability matrix defined by C = [B, AB, A2B...., 1 is full rank (Kalman 1960 Jurdjevic 1997). We instead evaluate the intrinsic dimension of the locally attainable set under feedback policies within our theoretical framework. We do so for the
set of states attained for small /. under the dynamics s(l) As(() + (+) where A, B are in equation 1T To achieve this, for fixed emhedding dimension d, obtain neural networks sampled uniformly randomly from the family of linearised neural networks as in definition 3 with 1.0, E (0, 5) 1024. Consequently, we obtain 1000 policies with öt 0.01, and therefore a sample of 500000 states to estimate the intrinsic dimension of the attained set of states using the algorithm of Facco al. (207)). We vary the dimensionality of the state space, d,, from 3 to 10 to observe how the intrinsic dimension of the attained set of states varies with the embedding dimension while keeping d, fixed at The dimensionality upper-bounded by 2d, = 3 for this system (figure This bound is even lower (at d, + 2) for linear environments because the Lie series expansion (equation 4 gets truncated GeLU activation owing to the fact that the second derivative close to zero in most of R.
the attained set states remains
for
5.4 REINFORCEMENT LEARNING WITH LOCAL NOW-DIMENSIONAL SUBSPACES
To demonstrate the applicability of our theoretical result, we apply a fully-connected sparisifica- tion ML.P layer introduced by Yu et al. 2023a) In series of works named the CRATE frame-

---

Published as a conference paper at ICLR 2025
Art
Dog Bolk
Quadiaaed Poakk
10:00
1300
0003
120
4003
400
2003
1003
230
230
Sumplen
Samples
Speples
1000 2008 2004
4832
2M
200c
100
(a) Ant
(b) Dog Stand
(c) Dog Walk
(d) Quadruped Walk
Figure 5: Discounted returns of SAC (blue) and sparse SAC (red) Or, = -0.5, arg = 0.4
work (Coding RAte reduction TransformEr (Chan et al. 2022 Yu et al., 2023a by Pai et al. (024). researchers have argued for better design of neural networks that compress and transform high- dimensional data given that is sampled from low-dimensional manifolds. They assumed that the data lie near a union of low-dimensional manifolds U,M, where each manifold has dimen- sion d, < dg. An innovation that has remarkable empirical and theoretical results under the manifold hypothesis learns sparse high-dimensional representations of the data Φ IRd -> IR" with n. 2 E, di. These representations are orthogonal for data points across two manifolds, Ty E M, E Mj.i # j
Φ(s.) ((s;) = O, and low-rank on or near the same man- ifold, rank((((=))) ((x)),.... ((,")) Rs d, for si E M,, This can be viewed as disentangling representations across different manifolds via sparsification. The sparsification layer of width n Yu ct al. 2023a) is defined as
z ReLU (z' +aWT (z--wz)) adi); -a>i)
(12)
where a is the sparse rate step size parameter, Z² is the input to the f-th layer, W are the n weight matrix. For further explanation of this sparsification layer, refer to Appendix M As is evident, this is a linear transformation of the feed-forward layer and therefore does not add computational over- head. To verify the efficacy of disentangled low-dimensional representations, under the manifold hypothesis, we replace one feed-forward layer of all the policy and Q networks with sparsifica- tion layer within the SAC framework (see Appendix K foror background on SAC). We also use wider networks of width 1024, for both the baseline and modified architecture, for comparison. This is to satisfy the assumption 2 E; d, described above. With simple code change of about 5 lines with same number of parameters and two additional hyperparameters, ax aq., we see improvements in the discounted returns for high-dimensional control environments: Ant (Brockman et al.) 2016). Dog Stand, Dog Walk and Quadruped Walk (Tunyasuvunakool et al., 2020), averaged across 16 seeds. Discounted returns are reported on the y-axis against the number of samples on the x-axis in Figure We observe that SAC with fully connected network fails to learn high-dimensional Dog environments where as SAC equipped with single sparsification layer instead of a fully connected layer does far better. This demonstrates the efficacy of learning local dimeesiioii representa- tions which arise from wide neural nets. We use the same hyperparameter for learning rates and entropy regularization for both the sparse SAC and vanilla SAC as those provided in the CleanRL library (Huang et al. (022). We report ablation for Ant and Humanoid domains over the step size parameter in Appendix
6 DISCUSSION
We have proved that locally there exists a low-dimensional structure to the continuous time trajecto- ries of policies learned using a semi- graiiegr ascent method. We develop a theoretical model where both transition dynamics and training dynamics are continuous-time. Ours is not only the first result of its kind, but we also introduce new mathematical models for the study of RL. In addition, we exploit this low-dimensional structure for efficient RL in nigh-dimensional environments with min- imal changes. For detailed related work, refer to Appendix O and address the broader applicability of our theoretical work in Appendix P We also assume access to the true value function Q, this is not practical and warrants an extension to the setting where this function is noisy. A key challenge that remains is extending this theory very high-dimensional datasets where d, - oo as n + We anticipate that noise in this settings will further complicate analysis. Additionally, the impact of stochastic transitions remains unexplored, our current analysis assumes deterministic transitions.
10

---

Published as a conference paper at ICLR 2025
7 ACKNOWLEDGMENTS
This research was partially funded by the ONR under the REPRISM MURI N000142412603 and ONR #N00014-22-1-25922 Partial funding for this work was provided by The Robotics and AI Institute, for which we are very grateful. This research was conducted using computational resources and services at the Center for Computation and Visualization, Brown University. We thank Sam Lobel, Rafael Rodriguez Sanchez, Akhil Bagaria, and Tejas Kotwal for fruitful conversations and guidance on the project. This work would not have been possible without their input. We would like to thank the members of Brown Intelligent Robot Lab for their continued support. We thank Wasiwasi Mgonzo, Alessio Mazzetto, Ji Won Chung. Renato Amado. Pedro Lopes De Almeida, Kalaiyarasan Arumugam, the second floor community the Center for Information Technology. and the larger Brown University graduate student community for their moral support and encouragement. The anonymous reviewers who have contributed very significantly to work over time have sincere gratitude.
REFERENCES
Ben Adlam and Jeffrey Pennington. The neural tangent kernel in high dimensions: Triple descent and a multi-scale theory of generalization. In International Conference on Machine Learning, pp- 74-84. PMLR, 2020.
Andrei A. Agrachev and Yu. L. Sachkov. Control theory from the geometrie viewpoint. 2004.
Anayo K. Akametalu, Shahab Kaynama, Jaime Fernández Fisac, Melanie Nicole Zeilinger, Jeremy H. Gillula, and Claire J. Tomlin. Reachability safe learning with gaussian pro- cesses. 53rd IEEE Conference on Decision and Control, pp. 1424 1431, 2014.
Zeyuan Allen-Zhu, Yuanzhi Li, and Yingyu Liang. Learning and generalization in overparame- terized neural networks, going beyond two layers. Advances in neural information processing systems, 32, 9a.
Zeyuan Allen-Zhu, Yuanzhi Li, and Zhao Song. A convergence theory for deep learning via over- parameterization. International conference on machine learning. pp. 242-252. PMLR, 2019b.
Rika Antonova, Maksim Maydanskiy, Danica Kragic, Sam Devlin, and Katja Hofmann. Ana- lytic manifold learning: Unifying and evaluating representations for continuous control. ArXiv, abs/2006.08716 18, 2020.
Ershad Banijamali, Rui Shu, Mohammad Ghavamzadeh, Hung Hai Bui, and Ali Ghodsi. Robust locally-linear controllable embedding. AISTATS, 2018.
Ronen Basri and David W. Jacobs. Efficient representation of low-dimensional manifolds using deep networks. ArXiv, 102/1400. 2017.
Gerard Ben Arous, Reza Gheissari, and Aukosh Jagannath. High-dimensional limit theorems for sgd: Effective dynamics and critical scaling. Advances in Neural Information Processing Systems, 35:25349-2553 253622, 2022.
Dimitri Bertsekas. Dynamic programming and optimal control: Volume 1, volume 4. Athena scien- tific, 2012.
Dimitri P Bertsekas. Model predictive control and reinforcement learning: A unified framework based on dynamic programming. arXiv preprint 240.00990 2024.
Anthony M Bloch and Anthony M Bloch. An introduction aspects of geometric control theory. Nonholonomic mechanics and control, pp. 199-233, 2015.
William M Boothby. An introduction to differentiable manifolds and Riemannian geometry. Aca- demic press, 1986.
Michael Bowling, Ali Ghodsi, and Dana F. Wilkinson. Action respecting embedding. Proceedings of the 22nd international conference on Machine learning, 2005.

---

Published as a conference paper at ICLR 2025
Roger W. Brockett. Lie theory and control systems defined on spheres. Siam Journal on Applied Mathematics, 25:213-225, 225, 1973.
Greg Brockman, Vicki Cheung, Ludwig Pettersson, Jonas Schneider, John Schulman, Jie Tang, and Wojciech Zaremba. Openai gym. ArXiv, (bs//1006.01500, 2016.
Michael M. Bronstein, Joan Bruna, Taco Cohen, and Petar Velivekovi'c. Geometric deep learning: Grids, groups, graphs, geodesics, and gauges. ArXiv, 15/2104,11478, 2021.
Sam Buchanan, Dar Gilboa, and John Wright. Deep networks and the multiple manifold problem. ArXiv, abs/2008. 1 1245, 2021.
Francesco Bullo and Andrew D Lewis. Geometric control of mechanical systems: modeling, analy- sis, and design for simple mechanical control systems, volume 49. Springer, 2019.
Keith Bush and Joelle Pineau. Manifold embeddings for model-based reinforcement learning under partial observability. In NIPS, 2009.
Q Cai, Zhuoran Yang, Jason Lee, and Zhaoran Wang. Neural temporal-difference and q-learning provably converge global optima. arXiv: Learning, 2019a.
Qi Cai, Zhuoran Yang, Jason D. Lee, and Zhaoran Wang. Neural immoral-diffeerenc learning verges to global optima. Advances in Neural Information Processing Systems, 32, 2019b. ISSN 1049-5258. Publisher Copyright: 2019 Neural information processing systems foundation. All rights reserved.; 33rd Annual Conference on Neural Information Processing Systems, NeurIPS 2019 Conference date: 18-1222019 Through 14-12-2019.
G. Carlsson, T. Ishkhanov, V. D. Silva, and A. Zomorodian. the local behavior of spaces of natural images. International Journal of Computer Vision, 76:1-12, 2007.
Pablo Samuel Castro and Doina Precup. Using bisimulation for policy transfer in mdps. In AAAI Conference on Artificial Intelligence, 2010.
Kwan Ho Ryan Chan, Yaodong Yu, Chong You, Haozhi Qi. John Wright, and Yi Ma. Redunet: A white-box deep network from the principle of maximizing rate reduction. Journal of machine learning research, 23(114)11-133 2022.
Minshuo Chen, Haoming Jiang, Wenjing Liao, and Tuo Zhao. Efficient approximation of deep relu networks for functions on low dimensional manifolds. ArXiv, htb//1988,0182,, 2019.
Daizhan Cheng, Xiaoming Hu, and Tielong Shen. Analysis and design of nonlinear control systems. In Analysis and Design of Nonlinear Control Systems, 2011.
Xiang Cheng, Dong Yin, Peter Bartlett, and Michael Jordan. Stochastic gradient and langevin processes. In International Conference Machine Learning, pp. 1810-1819. PMLR, 2020.
Lenaic Chizat and Francis Bach. On the global convergence of gradient descent for over- parameterized models using optimal transport. Advances in neural information processing sys- tems, 31, 2018
Alexander Cloninger and Timo Klock. Relu nets adapt to intrinsic dimensionality beyond the target domain. ArXiv, habs/2008.02545 2020.
G. Cybenko. Approximation by superpositions of sigmoidal function. Mathematics of Control, Signals and Systems, 2:303-314, 1989.
Mare Peter Deisenroth and Carl Edward Rasmussen. Pilco: A masel-based and data-efficient ap- proach to policy search. In ICML, 2011.
Kenji Doya. Reinforcement learning in continuous time and space. Newral Computation, 12:219- 245, 2000a.
Kenji Doya. Reinforcement learning in continuous time and space. Neural computation, 12(1): 219-245, 2000b.
12

---

Published as a conference paper at ICLR 2025
Simon Shaolei Du, J. Lee, Haochuan Li, Liwei Wang, and Xiyu Zhai. Gradient descent finds global minima of deep neural networks. ArXiv, abs/181 1.03804, 2018.
Yan Duan, Xi Chen, Rein Houthooft, John Schulman, and P. Abbeel. Benchmarking deep reinforce- ment learning for continuous control. In ICML, 2016.
Elena Facco, Maria d' Errico, Alex Rodriguez, and Alessandro Laio. Estimating the intrinsic dimen-
sion of datasets by a minimal neighborhood information. Scientific Reports, 7, 2017.
C. Fefferman, S. Mitter, and Hariharan Narayanan. Testing the manifold hypothesis. arXiv: Statis- fies Theory, 2013.
Norm Ferns, P. Panangaden, and Doina Precup. Metries for finite markov decision processes. In AAAI Conference on Artificial Intelligence, 2004.
Stanislav Fort, Gintare Karolina Dziugaite, Mansheej Paul, Sepideh Kharaghani, Daniel M Roy, and Surya Ganguli. Deep learning versus kernel learning: an empirical study of loss landscape geometry and the time evolution of the neural tangent kernel. Advances in Neural Information Processing Systems, 33:5850-555 5861, 2020.
Ruiqi Gao, Tianle Cai, Haochuan Li, Cho-Jui Hsieh, Liwei Wang, and Lee. Convergence of adver- sarial training in overparametrized neural networks. In Neural Information Processing Systems, 2019.
Carles Gelada, Saurabh Kumar, Jacob Buckman, Ofir Nachum, and Mare G. Bellemare. Deepmdp: Learning continuous latent space models for representation learning. ArXiv, abs/1906.02736, 2019.
Robert Givan, Thomas L. Dean, and Matthew Greig. Equivalence notions and model minimization in markov decision processes. Artif. Intell., 147:163-22 223, 2003.
Sebastian Goldt, Mare Mézard, Florent Krzakala, and Lenka Zdeborová. Modelling the influence of data structure on learning in neural networks. ArXiv, abbs/1909.11595 2020.
Anirudh Goyal, Riashat Islam, Daniel Strouse, Zafarali Ahmed, Matthew M. Botvinick, H. Larochelle, Sergey Levine, and Yoshua Bengio. Infobot: Transfer and exploration via the information bottleneck. ArXiv, abs/1901.10902 2019a.
Anirudh Goyal, Shagun Sodhani, Jonathan Binas, Xue Bin Peng, Sergey Levine, and Yoshua Ben- gio. Reinforcement learning with competitive ensembles of information-constraiii primitives. ArXiv, abs/1906. 10667, 2019b.
Anirudh Goyal, Yoshua Bengio, Matthew M. Botvinick, and Sergey Levine. The variational band- width bottleneck: Stochastic evaluation on an information budget. ArXiv, abs/2004.11935, 2020.
Victor Guillemin and Alan Pollack. Differential Topology. Prentice-Han Hall, 1974.
Tuomas Haarnoja, Haoran Tang, Abbeel, and Sergey Levine. Reinforcement learning with deep energy-based policies. In ICML, 2017.
Tuomas Haarnoja, Aurick Zhou, P. Abbeel, and Sergey Levine. Soft actor-critic: Off-policy maxi- mum entropy deep reinforcement learning with a stochastie actor. ICMI, 2018.
Boris Hanin and Mihai Nica. Finite depth and width corrections to the neural tangent kernel. arXiv preprint arXiv: met.:.999999999 2019.
Dan Hendrycks and Kevin Gimpel. Gaussian enror linear units (gelus). arXiv: Learning, 2016.
Shengyi Huang, Rousslan Fernand Julien Dossa, Chang Ye, Jeff Braga, Dipam Chakraborty, Ki- nal Mchta, and JoAGo GM AraAšjo. Cleanrl: High-quality single-file implementations of deep reinforcement learning algorithms. Journal of Machine Learning Research, 23(274): 2022.
David Isele. Alireza Nakhaei, and Kikuo Fujimura. Safe reinforcement learning autonomous vehicles. 2018 IEEE/RSJ International Conference on Intelligent Robots and Systems (IROS), pp. 1-6, 2018.
13

---

Published as a conference paper at ICLR 2025
Alberto Isidori. Nonlinear control systems: introduction. Springer, 1985.
Riashat Islam, Hongyu Zang, Manan Tomar, Aniket Didolkar, Md. Mofijul Islam, Samin Yeasar Arnob, Tariq lqbal, Xin Li, Anirudh Goyal, Nicolas Manfred Otto Heess, and Alex Lamb. Repre- sentation learning in deep rl via discrete information bottleneck. International Conference on Artificial Intelligence and Statistics, 2022.
Arthur Jacot, F. Gabriel, and Hongler. Neural tangent kernel: Convergence and generalization in neural networks. In NeurIPS, 2018.
Odest Chadwicke Jenkins and Maja Mataric. spatio-temporal extension isomap nonlinear dimension reduction. In ICML '04, 2004.
Xiang Ji, Minshuo Chen, Mengdi Wang, and Tuo Zhao. Sample complexity of nonparametric off- policy evaluation on oiw-dimesioonl manifolds using deep networks. ArXiv, abs/2206.02887,, 2022.
Yanwei Jia and Xun Yu Zhou. Policy gradient and actor-critic learning in continuous time and space: Theory and algorithms. Journal of Machine Learning Research, 23(25));1-50, 2022a.
Yanwei Jia and Xun Yu Zhou. Policy gradient and actor-critic learning in continuous time and space: Theory and algorithms. Journal of Machine Learning Research, 23(275);1-50, 2022b. http:http:/iit wwtp:/oopppppppap..aa.a.a
1387. html
Yanwei Jia and Xun Yu Zhou. q-learnin in continuous time. Journal of Machine Learning Research, 24(161):1-61, 2023. URLhttp::// /eaett://jmm..... wttpa//////pappee////2222---2-- html
Yuu Jinnai, Jee Won Park, Marlos C. Machado, and George Dimitri Konidaris. Exploration in reinforcement learning with deep covering options. In ICLR, 2020.
Velimir Jurdjevic. Geometric control theory. Cambridge university press, 1997.
Sham M. Kakade. A natural policy gradient. In NIPS, 2001.
Rudolf Kalman, On the general theory of control systems. In Proceedings First International Conference Automatic Control, Moscow, USSR, pp. 481-492, 1960.
Ioannis Karatzas and Steven Shreve. Brownian motion and stochastic calculus, volume 113. springer, 2014.
Elia Kaufmann, Leonard Bauersfeld, Antonio Loquercio, Matthias Müller, Vladlen Koltun, and Davide Scaramuzza. Champion-level drone racing using deep reinforcement learning. Nature, 620(7976):982-988 987, 2023.
Vijay Konda and John Tsitsiklis. Actor-critic algorithms. Advances in neural information processing systems, 12, 1999.
Charline Le Lan, Mare G. Bellemare, and Pablo Samuel Castro. Metrics and continuity in reinforce- ment learning. In AAAI Conference on Artificial Intelligence, 2021.
Jachoon Lee, Yasaman Bahri, Roman Novak, Samuel S Schoenholz, Jeffrey Pennington, and Jascha Sohl-Dickstein. Deep neural networks as gaussian processes. arXiv preprint arXiv:1711.00 101.0:17...0065, 2017.
Jaehoon Lee, Lechao Xiao, Samuel Schoenholz, Yasaman Bahri, Roman Novak, Jascha Sohl- Dickstein, and Jeffrey Penmington. Wide neural networks of any depth evolve as linear models under gradient descent. Advances in neural information processing systems, 32, 2019.
Sergey Levine, Chelsea Finn, Trevor Darrell, and P. Abbeel. End-to-end training of deep visuomotor policies, ArXiv, abs/1504/ 504.00004. 2016.
Yuanzhi Li and Yang Yuan. Convergence analysis two-layer neural networks with relu activation. In NIPS, 2017.
14

---

Published as a conference paper at ICLR 2025
Timothy P. Lillicrap, Jonathan J. Hunt, Alexander Pritzel, Nicolas Manfred Otto Heess, Tom Erez, Yuval Tassa, David Silver, and Daan Wierstra. Continuous control with deep reinforcement learn- ing. CoRR, 0bs//1509.299) 2016.
TP Lillicrap. Continuous control with deep reinforcement
learning.
arXiv
preprint
arXiv:1509.0297/ 2015.
Puze Liu, Davide Tateo, Haitham Bou-Ammar, and Jan Peters. Robot reinforcement learning on the constraint manifold. In CoRL, 2021.
Winfried Lohmiller and Jean-Jacques E. Slotine. On contraction analysis for non-linear systems. Autom., 34:683-696, 1998.
Yi Ma, Harm Derksen, Wei Hong, and John Wright. Segmentation of multivariate mixed data via lossy data coding and compression. IEEE transactions on pattern analysis and machine intelligence,
2007.
Marlos C. Machado, Marc G. Bellemare, and Michael Bowling. A laplacian framework for option
discovery in reinforcement learning. ArXiv, abs/1703.00956, 2017.
Marlos C. Machado, Clemens Rosenbaum, Xiaoxiao Guo, Miao Liu, Gerald Tesauro, and Mur- ray Campbell. Eigenoption discovery through the deep successor representation. ArXiv, abs/1710.1 1089, 2018.
Sridhar Mahadevan.
Proto-value functions: developmental reinforcement learning. Proceedings of the 22nd international conference on Machine learning, 2005.
Sridhar Mahadevan and Mauro Maggioni. Proto-value functions: A laplacian framework for learn- ing representation and control in markov decision processes. Mach. Learn. Res., 8:2169-2231, 2007.
Peter Marbach and John N Tsitsiklis. Approximate gradient methods in policy-space optimization of markov reward processes. Discrete Event Dynamic Systems, 13:111-148, 2003.
Song Mei, Andrea Montanari, and Phan-Minh Nguyen. A mean field view of the landscape two- layer neural networks. Proceedings of the National Academy of Sciences of the United States of America, 115:E7665 - E767 1, 2018a.
Song Mei, Andrea Montanari, and Phan-Minn Minh Nguyen. A mean field view of the landscape of two- layer neural networks. Proceedings of the National Academy Sciences, 116635)177665-777771, 2018b.
Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Andrei A. Rusu, Joel Veness, Marc G. Belle- mare, Alex Graves, Martin A. Riedmiller, Andreas Fidjeland, Georg Ostrovski, Stig Petersen, Charles Beattie, Amir Sadik, Ioannis Antonoglou, Helen King, Dharshan Kumaran, Daan Wier- stra, Shane Legg, and Demis Hassabis. Human-level control through deep reinforcement learning- Nature, 51 518(7540):529-533 2015.
Richard M Murray and John Edmond Hauser. case study in approximate linearization: The acrobat example. Electronics Research Laboratory, College of Engineering, University of ..., 1991.
V. Nair and Geoffrey E. Hinton. Rectified linear units improve restricted boltzmann machines. In ICML, 2010.
Henk Nijmeijer and Arjan van der Schaft. Non-linear dynamical control systems. 1990.
Jonas Paccolat, Leonardo Petrini, Mario Geiger, Kevin Tyloo, and Matthieu Wyart. Geometric compression of invariant manifolds in neural networks. Journal of Statistical Mechanics: Theory and Experiment, 2021, 2020.
Druv Pai. Sam Buchanan, Ziyang Wu, Yaodong Yu, and Yi Ma. Masked completion via structured diffusion with white-box transformers. In The Tweifth International Conference on Learning Representations, 2024.
15

---

Published as a conference paper at ICLR 2025
Gautam Pai, Ronen Talmon, Alexander M. Bronstein, and Ron Kimmel. Dimal: Deep isometric manifold learning using sparse geodesic sampling. 2019 IEEE Winter Conference on Applications of Computer Vision ( (WACV), pp. 819 828, 2019.
Quang-Cuong Pham, Nicolas Tabareau, and Jean-Jacques Slotine. A contraction theory approach to stochastic incremental stability. IEEE Transactions on Automatic Control, 54(4):816-88) 820, 2009. doi: 10. N0.11000AAC.20082000000 200969..
Balaraman Ravindran and Andrew Barto. Symmetries and model minimization in markov deci- sion processes. 2001.
Balaraman Ravindran and Andrew G. Barto. Model minimization in hierarchical reinforcement learning. In Symposium on Abstraction, Reformulation and Approximation, 2002.
Balaraman Ravindran and Andrew G. Barto. Smdp homomorphisms: An algebraic approach to abstraction in semi-markov decision processes. In International Joint Conference on Artificial Intelligence, 2003.
Jerzy Respondek. Controllability of dynamical systems with constraints. Systems & Control Letters, 54(4))293-314, 2005.
Joel W. Robbin, Uw Madison, and Dietmar A. Salamon. INTRODUCTION TO DIFFERENTIAL GEOMETRY. Preprint, 2011.
Johannes Schmidt-Hieber. Deep relu network approximation of functions on a manifold. ArXiv, abs/1908.00695, 2019.
John Schulman, S. Levine, P. Abbeel, Michael Jordan, and P. Moritz. Trust region policy opti- mization. ArXiv, abs/1502.004777 2015.
John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. ArXiv, abs/1 0an//707,006341 063406 2017.
Uri Shaham, Alexander Cloninger, and Ronald Coifman. Provable approximation properties for deep neural networks. ArXiv, 0tos/1509.073855. 2015.
Yifei Shao, Chao Chen, Shreyas Kousik, and Ram Vasudevan. Reachability-based trajectory safe- guard (rts): A safe and fast reinforcement learning safety layer for continuous control. IEEE Robotics and Automation Letters, 6:3663-3670, 2020.
David Silver, Aja Huang, Chris J. Maddison, Arthur Guez, Sifre, George van den Driess- che, Julian Schrittwieser, loannis Antonoglou, Vedavyas Panneershelvam, Marc Lanctot, Sander Dieleman, Dominik Grewe, John Nham, Nal Kalchbrenner, Ilya Sutskever, Timothy Lillicrap, Madeleine Leach, Koray Kavukcuoglu, Thore Graepel, and Demis Hassabis. Mastering the game of go with deep neural networks and tree search. Nature, 529:484-4989 2016.
William D Smart and L Pack Kaelbling. Effective reinforcement learning for mobile robots. In Proceedings 2002 IEEE International Conference Robotics and Automation (Cat. No. 02CH37292), volume 4, pp. 3404-3410. IEEE, 2002.
Yimin Sun. Necessary and sufficient condition for global controllability of planar affine nonlinear systems. IEEE transactions on automatic control, 528))1444-1166) 2007.
Zhendong Sun, Shuzhi Sam Ge, and Tong Heng Lee. Controllability and reachability criteria for switched linear systems. Automatica, 38(5):775-78)) 2002.
Héctor Sussmann. Orbits of families of vector fields and integrability of distributions. Transactions of the American Mathematical Society, 180: 180:171-188, 1973.
Hector J Sussmann. A general theorem on local controllability. SIAM Journal on Control and Optimization, 25(1):158-194, 1987.
Richard Sutton and Andrew Barto. Introduction to reinforcement learning, volume 135. MIT press Cambridge, 1998.
16

---

Published as a conference paper at ICLR 2025
Richard S. Sutton, David McAllester, Satinder Singh, and Yishay Mansour. Policy gradient methods for reinforcement learning with function approximation. In Proceedings of the 12th International Conference Neural Information Processing Systems, NIPS NIPS 99, pp. 1057-1063, 1999.
Russ Tedrake. mit.edu
Underactuated Robotics. 2023. URL httpp::/uunderactuatee.. csail.
Joshua Tenenbaum. Mapping a manifold of perceptual observations. In NIPS, 1997.
Saket Tiwari and George Konidaris. Effects of data geometry in early deep learning. Advances in Neural Information Processing Systems, 35:30099 3011001111 13, 2022.
Emanuel Todorov, Tom Erez, and Yuval Tassa Mujoco: A physics engine for model-based con- trol. 2012 IEEE/RSJ International Conference on Intelligent Robots and Systems, pp. 5026-5033, 2012.
Touchette and Lloyd. Information-theorriic limits of control. Physical review letters, 84 6: 6:1156-9, 1999.
Hugo Touchette and Seth Lloyd. Information-theoreti approach to the study of control systems. Physica A-statistical Mechanics and Its Applications, 331:140-172, 2001.
Saran Tunyasuvunakool. Alistair Muldal, Yotam Doron, Siqi Liu, Steven Bohez, Josh Merel, Tom Erez, Timothy Lillicrap, Nicolas Heess, and Yuval Tassa. dm_control: Software and tasks for continuous control. Software Impacts, 6:100022, 2020.
George E. Uhlenbeck and Leonard Salomon Ornstein. On the theory the brownian motion. Phys- ical Review, 36:823-841, 1930.
Elise van der Pol, Daniel E. Worrall, Herke van Hoof, Frans Oliehoek, and Max Welling. Mdp homomorphic networks: Group symmetries in reinforcement learning. ArXiv, abs/2006.16908, 2020.
Roman Vershynin. High-dimensional probability: An introduction with applications in data science, volume 47. Cambridge university press, 2018
Haoran Wang and Xun Yu Zhou, Continuous-ting mean-variance portfolio selection: A reinforce- ment learning framework. Mathematical Finance, 30(4):1273-1308, 2020.
Haoran Wang, Thaleia Zariphopoulou, and Xun Yu Zhou. Reinforcement learning in continuous time and space: A stochastic control approach. J. Mach. Learn. Res.,
2020.
Lingxiao Wang. Qi Cai, Zhuoran Yang, and Zhaoran Wang. Neural policy gradient methods: Global optimality and rates of convergence. ArXiv, abs/ htss/1999.0111 1550, 2019.
Sifan Wang, Xinling Yu, and Paris Perdikaris. When and why pinns fail to train: A neural tangent kernel perspective. Journal of Computational Physics, 449:110768, 2022.
Colin Wei, Jason D Lee, Qiang Liu, and Tengyu Ma. Regularization matters: Generalization and optimization of neural nets vs their induced kernel. Advances in Neural Information Processing Systems, 32, 2019.
Stephen Wiggins. Introduction to applied nonlinear dynamical systems and chaos. In SPRINGER, 1989.
Ronald J Williams. Simple statistical
algorithms for connectionist reinforcement
learning. Machine learning, 8:229-256. 1992.
John Wright and Yi Ma. High-DDmensional Data Analysis with Low-Dimensional Models: Princi- ples, Computation, and Applications. Cambridge University Press, 2022.
Yifan Wu, G. Tucker, and Ofir Nachum. The laplacian in rl: Learning representations with efficient approximations. ArXiv, abs/1810.0455e. 2019.
17

---

Published as a conference paper at ICLR 2025
Peter R. Wurman, Samuel Barrett, Kenta Kawamoto, James MacGlashan, Kaushik Subramanian, Thomas J. Walsh, Roberto Capobianco, Alisa Devlic, Franziska Eckert, Florian Fuchs, Leilani Gilpin, Piyush Khandelwal, Varun Kompella, HaoChih Lin, Patrick MacAlpine, Declan Oller, Takuma Seno, Craig Sherstan, Michael D. Thomure, Houmehr Aghabozorgi. Leon Barrett, Rory Douglas, Dion Whitehead, Peter Dürr, Peter Stone, Michael Spranger, and Hiroaki Kitano. Out- racing champion gran turismo drivers with deep reinforcement learning. Nature, 602:223 - 228, 2022.
Huaqing Xiong, Tengyu Xu, Lin Zhao, Yinghin Liang, and Wei Zhang. Deterministic policy gra- dient: Convergence analysis. In Uncertainty in Artificial Intelligence, pp. 2159-2169. PMLR, 2022.
Greg Yang and Edward J Hu. Tensor programs iv: Feature learning in infinite-width neural networks. In International Conference on Machine Learning, pp. 11727-11737. PMLR, 2021.
Greg Yang, Edward J Hu, Igor Babuschkin, Szymon Sidor, Xiaodong Liu, David Farhi, Nick Ry- der, Jakub Pachocki, Weizhu Chen, and Jianfeng Gao. Tensor programs v: Tuning large neural networks via zero-shot hyperparameter transfer. arXiv preprint arXiv:2203.034660 2022.
Yaodong Yu, Sam Buchanan, Druv Pai, Tianzhe Chu, Ziyang Wu, Shengbang Tong, Benjamin David Haeffele, and Yi Ma. White-box transformers via sparse rate reduction. In Thirty-sevntth Con- ference on Neural Information Processing Systems, 2023a.
Yaodong Yu, Tianzhe Chu, Shengbang Tong, Ziyang Wu, Druv Pai, Sam Buchanan, and Yi Ma. Emergence of segmentation with minimalistic white-box transformers. In Conference on Parsi- mony and Learning (Proceedings Track), 2023b.
Amy Zhang, Rowan Thomas McAllister, Roberto Calandra, Yarin Gal, and Sergey Levine. Learning invariant representations for reinforcement learning without reconstruction. ArXiv, 2020.
Henry Zhu, Justin Yu, Abhishek Gupta, Dhruv Shah, Kristian Hartikainen, Avi Singh, Vikash Ku- mar, and Sergey Levine. The ingredients of real-world robotic reinforcement learning. ArXiv, abs/2004.125700 2020.
18

---

Published as a conference paper at ICLR 2025
Appendix
A MANIFOLD BACKGROUND
Here, provide precise definitions which essential to the theory of differential geometry but might not be absolutely essential to understanding our results in the main body of our work. The tangent space characterises the geometry of the manifold and it is defined as follows.
Let M be m-manifold in Re and p E M be a fixed point. A vector E TR is called a tangent vector of M at p if there exists a smooth curve y + M such that 7(0) = 00)( = v. The set T,M = (+)) E M smooth, 70) p} of tangent vectors of M at p is called the langent space of M at p.
Definition 7.
Continuing our example, the tangent space of a point p in S is the vertical plane tangent to the cylinder at that point. For small enough and a vector E Tp.5² there exists unique curve y [-c,c] - 22 such that 7(0) p and +(0) v. The union of tangent spaces at all points is termed the tangent bundle and denoted by T(M). At point p. the tangent space 7,M is spanned by
{
0
the vectors|
Any tangent vector e E TyM can be expressed as a linear combination:
Oxt
Dx.
1p
where of E Th. are the components of y in the basis
Or. p
B FEEDBACK ACTION LIE SERIES
Consider the vector fields for feedback policy a(x) C
x +=g(x)++Mrx)
Consider the Lie series and its first term:
cx = + Σ 13 x (x)
Lx [z)=((b>>) + h(z)a(x))
The second order term be written as:
d,
0ge(x) + hx(2))(())
(T22 x)
(x) + Inzn(a)as(o)
Oz
1=1
1=1
d,
(x(x)) (g:(a) dok
+ da
(gi(a)
hz)(()aa((z)) drk
1=1
du
The (x)
ag (x) + (xy
g)(z))
dxx
dak
de
du
(h.
a (x) + ha. (2)())) Oay(x) axx
mos(
h (x))
Oxx
19

---

Published as a conference paper at ICLR 2025
While we do not use this third order term present it to demonstrate that tracking the statistics of this gets exponentially difficult:
d,
da (x(x) + Σ
(x
(L3, x); =
Σ
(2)(())(()))
==1
dg
dn
Oga(v)
+ 9k (x) (g)(x)
(c(T) + Σ (x)a
doe
Oox
1"=1
Oligh (x)
@g.(x)
(2)(2)) (gi(x)
(g)
05(x)
+ (z)((())
+ (),(())
dIw
drk
Ork
Ōrk
drr
(zx(z) She (z) (f((x) + gx(x)
(2) ay((x) + gx(z)
y day
Dow
Ozk
0zx
Ox
dik
(g)(z) (ay( (x) + (x)
Fay (I)
(x) + gu (x) )ay(2) (2) Bug.Oa hig (x)
Oaw
Oek
Oak
(h),)(T)
az(x) (a)) ay(x) + hz)(2) 0az(x) (lig((z) ay(x)) ()
Phag(1))
Oxy dor az(z)
Oxx
Ork
Oxw
Oxk
+ (2)(())a (x) (z) 0hz,a(z)
az(x) day((z) hay (x) + h2
daz(x) day(x)) h
drk
OIw
dow
Ozk
Oxw
drk
 ar( bid (x) ()n(())a (ay(1) (2)))()) ))
+ (z Oxr
Ook
Oxw
C
SOME HELPFUL DERIVATIONS
Here we derive various expressions to bound their magnitude terms the width of the NN n. First we consider the gradient term:
N VaQu (8; ai, (s;; Wo)
G(Y) lim No +00 N
lim
of (bs) Wo)
+ N'
d_ st) Σ
lim
sk
Co
N N' 71
J.
1=1
mid,- 11)
N²
lim
Σ s (wo. (()))((s))
Vn
J.m
No0000
No
11+
du
Σ
G2. Go
mld+-1)+
20

---

Published as a conference paper at ICLR 2025
The fourth equality suggests that individual elements of the vector G at any point are iid samples from the same random variable. Similarly we expand the term for (Y)
B
d
d
M Vn
{(Y)
BaY
L
CD J.m (W..
Co 3.
G j.mldd (1)
B
b=1
míd, -1)
1 EBY
de
B
da
Co
B
Σ
G 1mdd
n2
Ja
b=1
5=1
1,(d,-1)+4
d.
Σ
Co EBY
B (5((5
x p(WWn 557
Σ
G (,mldd-1)444711
Vn
míd. -1)
B de
du
MY = 7 BEa
C J.m &'
2j (s J0
Σ
c0
J.m j, (, -1)(dd)1))+(())
G
B
b=1Ø=1
B
da
de Gym Cm DC (Wa. 88 )2)
If (s
B
0'=1/==1
3'=1
mld. kk) 11
dg
B d
ERY
46 (5"))a (5" JCyms
&m (W (8 100cCmmmmm s' (Wor me )0%
n
B
0'=1/==1 1= j=1
de
da
Lj C d'mw Cm ma (8)
j=1
dg
B
da
CO G j.m(d Ja
bo (Y) )CO m s8 )
B
d,
do
Jim J.mGGmmd,-II4H (Y)C), Gr md (1)
1=
f==
midd 11)+k, (d, (d,
which we combine to form:
dg dg
M²
17
c
1
Σ
Σ
(w
(s) (w )5
71
B
==1===
Gm 4o (s CW/ s')sk m -1) 1)))
B
1=1
1
(3(50) (W . ")st&,mmle--utt
B
G G) 1) (Y)
(d.11)+
m(d, (d, 11+k
21

---

Published as a conference paper at ICLR 2025
Let the internal term in the summation be defined as follows:
B
B
H
-1, =EaY
1
Σ
st)shar (s (s*))'((())))))))))) (W'WWn st )sy
mld,- (ilds-1)+k,m'
B
b=1
V=
B
1
((5))Cu'(WW
(d.
B
b=1
1
b (s (o'() ( -0
B
+ G j,m(d, (m
(
B
=BBY
1
Σ
(W 8h yot -
W=1
(s >W'(Wn ss - ()mm((e--1))40(())
B
D COVARIATE TERMS
We denote V.Qu a, ta)la=as by (91 ....,........................... qa, (5*) shorthand. Consider the term:
E
M A: Av M,
= (s) ((()))
(B, Wo) (Y) Wo)
n
du
d,
33/2 (s, Wo)
Σ
Com Co mr m' (d.-ddd+aae 1)) (X)
H...'
G 88)8
m1)
22

---

Published as a conference paper at ICLR 2025
We further expand the dot product with the nd, 1 vector ss Wo)
d2
d
d
d.
E
M
M 3/
Amm Cp mm
ml Co j.
3n²
m=1 k= me
I=1
P=1
H1rl"
m'd, - +k.mml,, -1)+k (Y) (W me 8)SW (W 8)88
d
d
dg
da
1
ΣΣΣΣX
Cl. Bum Gp C2
3n m=1
s'(Wn, m'
s)swee' m 8)88
sh) (s (s Do'(WW me S
E E By
B
b=10=1
B
91(8 Cip (m (1200))))
1
B
q (e')''(WW,,, Gimt
B
b=
G Gmmdd 11)+ (Y (m (1)))))))))))))))
2
EBy
xe (W s)sw ssss
3n²
B
B
(C), (mmm)) (CO j.ml )²
( ( s5)s o Sh )sk
B²
V=1
(c j.m) )2
B
b=1
)
Σ
(8" (WW' 56 18% -1)
+ me (CO))) " j.m(d, ()m (Y)
jr,ndd (1)(s.)))++))))) (m,
2
M2
3n
I.' L''''''
In the n + do we note that 2 M.y... + because we have [[C|EE[Cr] + as a multiplicative term, while the other terms are finite and bounded in second moment because of the boundedness
23

---

Published as a conference paper at ICLR 2025
properties of gradient GeLU activation y For j = j' we have the following expression
B
M
MMM
M
2
BB
3 (W mi s)s 8)88
3n2
B
B
(C),er)*
15 (o'))((W' (1) (W 8 )ok
B
#=1
B
-(C J,m!
93 (1 Cip'(W &
(13)
B
hml
Jam)) )1
4b )o (W . G J.mid
B
+0 (0(H)
(C9
where the ( term result of the convergence rate of the strong of large numbers (Ver-
shynin 2018
E SUFFICIENT STATISTICS
We want to determine the dynamics of some linear or quadratic function of these parameters, for example the j -th output of the policy network Ay(s; W) (s Wo)W. Therefore, we find setting where continuous time SDE such as
dAy(s; () p(Ay(s; W))dr + (s) (s):::)
(14)
represent the dynamics of As Wny). In other words, find the conditions under which Wkyy Yon are close together in some sense.
Furthermore, assume that the activation, js, has bounded first and second derivatives almost everywhere in R. This assumption holds for GeLU activation. Moreover, we would like to show that we can track the statisties corresponding to elements of the Lie series
A7(s) (s; WT).
0A7(s) 47,x(s) = drk
(3)
(s) Oxx
The manner in which the push forward of the distribution of parameters described in Section e,
X XWW)) (s), changes with gradient steps is central to our work. In this Section we derive the suf- actions and their quadratic combinations
ficient statistics that determine how the distribution evolve over gradient steps. We present the following lemma for the learning setup described Sec- tion and the sufficient statistic required track the changes in the action A5(s; W)) as described
24

---

Published as a conference paper at ICLR 2025
in Appendix E First, define the following random variables for fixed E S:
W)
1
Σ
C s) wm).s
vn m=
C2 Cmm (W s) wi))
('Wn ss)AWm
amm L (W Σ 8y AW mld,-1)+k
G.mi (W
(mm(a 1)+t SW& (d,--144
()
where AW W W (s;; Wo) as defined in Section 3.1 and AWm, m,k Zz.as;; W) is the random variable that determines the action j corresponding to the action and state space dimension given parameters W. This is because individual random parameters AWm, &.m, can be viewed as i.i.d samples from distribution (see Section C
Similarly, we derive (*7(*) from the fifth equality above as follows:
(s) Oak
1 d,
-144 mn(d-11)+ 8)8k
j.m AW
Osk
m=1
where once again the i.idd copies of random variable: C2 AW m, Zj W) appear in the expression. Therefore, to track the random variables corresponding to the quantities of interest in the first and second order Lie series in Section B we need to track d,d, random variables: Zj,a with j E (1....,da);4 E (l....,dah.
F TRACKING STATISTICS
S denote that is less than y times some constant. We also write L% (ER) denoting the supremum of a function that depends on on the compact set K, As noted in Section E we aim track the following statisties for fixed across gradient steps
Notation: In this Section
(s, W). (s, W) (s, W), A,,, (s, W).
Moreover, we seek to derive their dynamics in the continuous- terisation of two-layer network policy (equation ( and the gradient update is as described in Sec- tion 3.2 We present lemma, whose proof follows the proof of Theorem 2.2 provided by Ben Arous al. (2022). except that in our case the dimensionality of the input data remains constant, the dynamics of summary statistics linear in the parameters that descrihe the learning dynamics under SGD. Here the sufficient statistics that determine the dynamics of the action A denoted by X,. We prove the dynamics of the j-th action.
limit. Given linearised parame-
Lemma 8. Given fixed state the j-th action, A7 determined by a linearised neural pol- icy with two hidden layers as described and initialised in Section 3 we assume Wo ~ Xn Normal(O, 0, Ta. /a,dd.) i.i.d whose gradient dynamics are described in equation 8 with learning rates Th + O, and under assumptions 4 we have that in the limit n + co the dynamics of A converge weakly the following random ODE
(15)
with the random variable X, is the limit point of the sufficient staistics, X of the parameters updated according stochastic policy gradient based updates laid out in Section 3.2 with 1l Tp--
25

---

Published as a conference paper at ICLR 2025
Proof.
Suppose the evolution of W over gradient steps is follows
W = =W + 17 (W, B.,,,), where
5n Bw_..
Σ
VoQ" (sb. a))== (sb, ab).
We further let G,, (W, -BB,,-, (5n (WW Let H (W, (By) = B(W Further let
(W EBw (W, (W ]]
For the statistic Aÿ consider the following evolution
vj - =
=
(16)
5f --1 1-1 = Hn ; Bw,),
where h represents the feature vector for the j-th action at state 8. Omitting the subscript in Yh since we take the limit 7/n -> 0 - oo, we consider the uj follows,
uj uf +
(1j oz
5 -1)
r'=1
T'=1
Now for l E [0, LJ define,
sj(1) -w//ni -.//1-1 and s,(1) s//10l -j
11/a1-1
If we let
17(00) = [. 15(1)a1' = (/((1l/1n)) + (s (s- ml(/nn) (u//ml -u 04()) [. <,(T)dt' <(((i1/1))) + -n(lni) (S//il - -51/n1-12)
be the continuous linear interpolations based on the discrete random variables v, and combine them together to obtain
(1) of (0) + و3 (1) + o1))..
(17)
Given a compact set K C IR and the exit time тx we aim to show that for all 0 s
E Ell ( of Tx) SK, (1 a)+
where the expectation over the stochastic updates. This proves that uf(ss A TK) is 1/4 Hölder- continuous by Kolmogorov's 's continuity thcorem (see Section 2.2 in the textbook by Karatzas & Shreve (2014)). We have for all s,t
[|:] of (1Il || - (1)l (s) 0011-
For a fixed W/ the action j corresponding to the linearised policy in the limit + oo is defined as: āz(8, W) = lim
== (s)
Now further suppose W is a stochastic variable where each of its entries is sampled i.i.d. from some distribution X e P()), where (R) is probability space over with the canonical sigma algebra. Therefore, the push forward of this stochastic random variable W,, in the limit n - can be defined as:
()) lim Σ
co
Am (Wo
8, W d, (m-1)+is
m=
26

---

Published as a conference paper at ICLR 2025
which converges in distribution to Normal distribution, with the mean and the variance are depen- dent on the state and the distribution Xn, by the Lindeberg-Lete central limit theorem (Vershynin 2018 Cai et al., 2019b) which is also consequence of using GeLU activation which has bounded derivatives a.e. We drop the argument X in the wherever it is implicit. The first term in the inequal- ity, the norm of ug,, upper-hounded as below:
W/lAAR in Σ
E - SKB
1
GGW(H)
r' <'=ls/TAATi//
$ " (z
SK. (t st=
where the second inequality is from the continuity the function D G(WT)) in W The last inequality is from the Lipschitz condition the gradient function G For the second term, which a martingale, in equation 17 we seek a similar bound to the one presented above:
[t/O)ARK In Σ
E -o](AA))) =E
(5 (5-
(s//////TT//
47
[t/o)/rk//1
=E
11
Σ
D Hn HW, Bw.,,,
+/ (1//////TK//
27
[t/]AAK///
$ E
"1
Σ
( (W, ))²
r'=(s/o]ArK//
where the last inequality is from the Burkholder's inequality. We further expand the last term in the inequality as follows:
[e/l/aTW// /n
E
Σ
( (W, Bw,)))
E
( H (W, BWW.B.,)) 1) (0) HW, Bw,,.))
T7 pp
1/2 2
Σ
2 (0)
SK.LD (t s)²
where the inequality in the second line is from Cauchy-Schwarz and for the last inequality we use the fact that (assumption/4) and the fact that 1 + 0 at rate 0 (,)) This proves that a is 1/4 Hölder- continuous by Kolmogorov's continuity thcorem. Since both the sequences μz and oz are uniformly 1/2 Hölder-comtinuous we have that vf (sAAk) (equation 17 is also 1/2 Hölder-continuous Further, we note that (s A TK) forms tight sequence in n with 1/2 Hölder-continin of (s A TK) - 145 TK) is martingale with a martingale limit point that 1/4 Hölder-continuous limit point.
limit point and
Now that we have proved that limit points exists and are 1/2 Hölder-continuous we seek to derive this limit. To do derive the quadratic variation
S.
(t A K22
which is a Martingale process. We seek to derive expression for the expectation above in the limit + To do so we derive the following:
En En
[(0]] H, (wpp)// Bw: )y.
(
27

---

Published as a conference paper at ICLR 2025
where omit the subscript BWijnat under the expectation in the expression right side for brevity. Letting Y WL/vIATK and writing out rhe above expression based on equation 13
d
(X))))
3n² B3
X p'(w s)sp (W s)sk
k= m
(Cimm )
2j ((8))/(W
3 (or Do'(Wn 5 )8%
B²
- -CC Jame )*
9 Cap' (W.)
B
b=l
1
(C) )
95(s (g' (W 5y G (2)46)-1))+(()))
B
+0 (0)
+ (mmaa
Given that the gradient updates have finite and bounded variance (assumption the expression including the expectation converges to a value that O(1) and dependent on 8, GY), H. by the strong law of large numbers at the rate O (2). .. We have therefore have the following
lim D (W)))) 0.
Therefore, as - oc and by localization technique (Karatzas & Shreve 2014) we prove conver- gence of equation t7:
ds, WT) v(s; Xt,
(18)
which admits unique solution due the assumption of Lipschitz condition ,assumption 5 Given
that we initialise Wn drawn from distribution in P(R) then the distribution of Aj push forward of this distribution and therefore evolves in equation T8 and gives us the result.
Similarly, from the linearity of AJ, in W using similar derivation as above we can derive an ODE. To do so we first note
A7, (s: W) 37 ( + Φ" (s)W;
where 1 is ddddiennsoonal vector with value at index is set to I and rest 0, 2 (s) is defined below
1 Vn
t (s)
[WD CY s)ST, W: Cagge""""" S)ST .... W C Ch s)sT] RIxnd.
Therefore, in the limit + xo
= u'
Now derive and prove the dynamics of the quadratic term A5(s; (s:W.,
Lemma 9. Given fixed state the j-th action, A7(x;) ,), determined by a linearised neural pol- icy with two hidden layers as described and initialised in Section 3 we assume Wo Xo = Normal(O, la, /d,) i.i.d whose gradient dynamics are described in equation g with learning rates Ta - O, and under assumptions 4 we have that in the limit n - oo the dynamics of AT converge weakly the following random ODE
(s X- ( Xx) (1y(5; X. (2
+ X=)
dr,
(19)
with the random variable X, is the limit point of the sufficient random variables, X of the pa- rameters updated according to stochastic policy gradient based updates laid out in Section 3.2 with
n Un. and Vj Vy are described in Lemma 8
28

---

Published as a conference paper at ICLR 2025
(s: (s; - ( (Q]
Proof. Since we know that A7, follow the ODE in equation 15 we want to show that the update
for AJ (s: W) A3 (s; W) converges weakly
d(A, (s; X,) (s (x)) = (u, ( X;) + vy (s; (s; dt
To do consider the increments in equation T6 for the statistic which product of two actions
Aス A3:
S3 455
5-1 (W,, by (9,W,,
Bw.) +n H (W,; Bw,)) 87 G.W,) + (n T7. Hn (W,, Bw,) ))
Omitting the subscript in 17n we obtain
u = ug + E [Vj vg -1
+ ($] sj 11
r'=1
Now for l E [0, L] define,
v/110) - 1/j -,
1//l 11/11-1 -5((1) =s ///11 s. 1//l111
Similar to the rpevious proof, we let
k²(0) = I. 17(0)d) = (/l(///))) + nl|/n|) (u///ol -
-1 -///n-11)
01(() / <((()d(' <,(iil/mm) (s n/(/ni)) (5 - - 11/11111)
be continuous linear interpolations based on discrete random variables v, and combine them to- gether to obtain
o) - +10)
(20)
With exit time we want to show that for all 0 s,tsT,
E (s (5 Tx)ll SK. (t s)²
where the expectation over the stochastic updates. This proves that uf A Tк) is 1/4 Hölder- continuous according to Kolmogorov's continuity theorem (as opposed to 1/2 in the previous proof).
We have for all 8,
(s)](s) - (0)l s || - x3(00 + (s0](() (t)l.
The first term in the inequality, the norm of uz. is upper-bounded as below:
WN/RAARC/T /n
E - Sx E
Σ
+
I (t (.-,)
+ W. (1~
SK.O.s (t - 5)
29

---

Published as a conference paper at ICLR 2025
where the second inequality is from the continuity of the function 0; W, 1 G,,,,,, in W". The last inequality is from the Lipschitz condition on the gradient function G
Further consider the second term
[t/nAaK//n In
E - =E
E
(5)
53'11 15--5'-11
"
T'
WiinAK// /n E
47
of (--5-11
E
"
[/=/s///TK///
W///ATK/1 /1
(5 (5
2 E
Σ
r'=ls/n|ATK/n
Using Cauchi-Schwarz inequality we can bound the two different types of terms separately. We can first upper bound
E
1/2
Σ
(Q)
7'=/S//AAK//
(,, D2 W,)) }
E
1,² ((0g (w) ()) 1/2) 2
03W,))))
11
S x)²
where the inequality in the second line is from Cauchy-Sehwarr and for the last inequality assumption 4. + 0 at rate I and the fact that b, only depends on s. Similarly, we bound the other term below
Vn
[t/)ArK//n
27
E (0)H)(W.,)) (3
E
(3/n)ArK/1
n
E
(10) H. (W+)) (W,))) (()) (W..), (W..))
T',""
2 ((*)H, (W)))))))) 1/2 2
SK.La (}=s)2
Where the last inequality follows from Assumptions 4 convergence of v) - 0 and Lipschitz conti- nuity of G. Combining these together we obtain the 1/2-Hölder-commimuuus limit point. Finally, to derive the dynamics we once again show that the quadratic variation goes to 0.
of ( - 1'' THE W Hn(w)) Dy,, + (9) D3 Hw,)))
Since we have already shown that Ex W [ (W [/
(W )){])))(0-00 = - 0(1/n) (see Sec-
tion similarly we have (Th w- = 0(1) + O(1/n). Therefore, as n - and by localization technique we have as required:
(A)
(v, (:) x2) + uy (x) xx) dt
30

---

Published as a conference paper at ICLR 2025
Similarly, can show the convergence to random ODE for the product of two linear terms Ag
G PROOF OF MAIN RESULT
Proof. We put together the dynamics we have derived and proved convergence in the -> oo limit. We re-arrange the terms of second order expansion of the Lie series to group together the terms that have the same multiplicative vector. We additionally denote the vector dg(e) by g (s) similarly Bsk
(o) as (s). From the derivation in section we have the following: Usk
x(w) (s) =s g(s) + (s(s)A)(s:W)
t2
gh (s(ggg((s) +
Σ han(s)A,(s,, W)g((s)
du
+ (s)
Σ
d
du
(L)())))(s.W)), W) (,
f=1
d,
=s =s + (tg(s) + Σ gx(s)g((s)) (s))
Σ (h)(8)
de
A)( +t
Agr + Σ (
Σ (x)
d
de
h)(x())A) ())
1=1
k=1
J'=1
2 da (y W) Σ 3'=1
de
(s)) Σ
gh(s)) + Σ h), W) 1=1
The first four expressions in the summation account for da + degrees of freedom. In other words, the first three terms in the summation are spanned by da + vectors. For the last term in the summation consider the following representation:
17. (s)
Σ hly (s) (1aa(6)))TT))
do Σ (7(())1(11())
-s](s) Σ My (()
d'=1 As] (s))sss)))
Where Jhy is the Jacobian of the function (s). This leads us the following vector:
de
(s)(s)
da
1 ask Σ
(s))
= (s)h(s))(s)BBB)
B3 represents the d, x vector [E[4]]]], B[4 Ai]] where the expetation is over the stochas- ticity of initialisation..
where (s) is the d, d, Jacobian of h, w.r.t s, aj (s) is the process determined by E [v]]. and h(())) concatenation of d, vectors. We seek upper bound the following, by showing that there
31

---

Published as a conference paper at ICLR 2025
is continuous time martingale process in the limit n 00.
DV" 0') = D
fj. Sp , v.,)),
==1
where D(f 6²) is the distance between f and the span of vf, ui..... v/ This distance is upper bounded as follows:
de D(ff.u)) .z))
DV .0) SI²
j=l
da
S
Σ 00)
du Σ
<
] BJ)
ST?
where both ā,, A are the mean processes, and L denotes the L-2 norm. Therefore, the data is concentrated around the space spanned by the vectors: hi...., hd, vi, ...., 22. and the paraboloid tg + l'g'. Let this product space be M then dim (M) I 2d, + 1. The concentration property is a result of the concentration of az (s) around B3 due to the dynamics in 19
While proofs Appendix Appendix(Fclosely cloosely follow that of Ben Arous al. (2022) we list our contibtuions in this work:
1. We show that the distribution of outputs, and its quadratic combinations, of two-layer linearised NNs deviate only in mean and variance, and dependent finite set of summary statistics, despite the width and parameter size going to infinity as the learning rate goes O.
2. In this appendix section, We combine this with the idea of the exponential map being push forward of the parameter distribution at gradient time step т, for fixed state s, and show that the distribution is concentrated around low-dimensional manifold.
H CONTINUOUS TIME POLICY GRADIENT
In this section we define the continuous time policy gradient. The policy gradient is similar to the one for discrete time introduced in Appendix ! We define the i-th component of the gradient over the value function with respect to the actions as follows:
(s hee, (" (s,
(V (s,
lim h-
where e; is d, dimensional vector with i and 0 otherwise, finally the function Qn is defined as:
Q7 (s, a + hea,t) = (8(+6) + 1- e"' fr(ppthei)d such that
s ghthe: he: =T(s, a he,
where T is the transition function defined in Section 2.1
This is application of the policy
gradient theorem (Sutton et al. 1999 Lillicrap 2015) which is a semi-gradient based optimisation
technique.
32

---

Published as a conference paper at ICLR 2025
It is not always guaranteed that the policy gradient algorithm will converge to globally optimal policies for general dynamics in a straightforward manner (Sutton et al. 1999, Konda & Tsitsiklis. 1999 Marbach & T sitsiklis 2003; Xiong et al. 2022). The rate of convergence affects the constant R2 introduced above. While we do not argue this formally, our main result holds as long as the RI. algorithm does not diverge.
I DIMENSIONALITY ESTIMATION
We describe the algorithm for dimensionality estimation in the context of sampled data from the state manifold S.. Let the dataset be randomly sampled points from a manifold Se embedded in Idd denoted by D { For point s; from the dataset D let {ri,, Ti...... } be a sorted list of distances of other points in the dataset from 8; and they set ro = O. Then the ratio of the two nearest neighbors is jr; r2,2/rs,1 where r1,1 is the distance to the nearest neighbor in 'D of and r5,2 is the distance to the second nearest neighbor. Facco et al. (2017) show that the logarithm of the probability distribution function of the ratio of the distances to two nearest neighbors is distributed inversely proportional to the degree of the intrinsic dimension of the data and we follow their algo- rithm for estimating the intrinsic dimensionality. We describe the methodology provided by Facco et al. (017)) in context of data sampled by an RL agent from a manifold. Without loss of generality, assume that {sx) N are in the ascending order of ri. We then fit line going through the origin } The slope of this line is then the empirical estimate of tim(S,)) We refer the reader to the supplementary material provided by Facco et al. (2017) for the theoretical justification of this estimation technique. The step by step algorithm is restated below.
for (log(44), log
1. Compute 1i,1 and ri22 for all data points i.
2. Compute the ratio the two nearest neighbors μ4: T2,2/T1.1
3. Without loss of generality, given that all the points in the dataset are sorted in ascending order of μ the empirical measure of edf i/N.
4. We then get the dataset Daensty {(log(4:) - log(1 i/N)) through which a straight line passing through the origin is fit.
The slope of the line fitted as above is then the estimate of the dimensionality of the manifold.
J DDPG BACKGROUND
An agent trained with the DDPG algorithm learns in the discrete time but with continuous states and actions. With abuse of notation, discrete time and continuous state and action MDP is defindd by the tuple M = A, P, fr, 8p. A), where S, A. so and fr are the state space, action space, start state and reward function as above. The transition function P : S x x S is the transition probability function, such that P(s, 8') Pr(S.+1 s'|S; s, A, a). is the probability of the agent transitioning from s to s' upon the application of action a for unit time. The policy, in this setting, is stochastic, meaning it defines probility distribution over the set of actions such that (8, a) = Pr(A, alS, = s). The discount factor is also discrete this setting such that an analogous state value function is defined as
T Ix as)ss
0= (8t) = Egga
which is the expected discounted return given that the agent takes action according to the policy т, transitions according to the discrete dynamics P and s; is the state the agent at time f. Note that this is discrete version of the value function defined in Equation
The objective then is to maximise J ((T) v. (su). One abstraction central to learning in this setting is that of the state-action
value function Q S S A R, for policy n. is defined by:
T
Qa = EB,
X (r as,,
33

---

Published as a conference paper at ICLR 2025
which is the expected discounted return given that the agent takes action a, state 8t and then fol- lows policy и for its decision making. An agent, trained using the DDPG algorithm, parametrises the policy and value functions with two deep neural networks. The policy, S - A, is parameterised S x AA - RRis also parame-
by a DNN with parameters 0 and the action value function, terised by DNN with ReLU activation with parameters 02. Although, the policy has additive noise, modeled by an OrnsteinUUllenbeck process (Uhlenbeck & Ornstein [130),, for exploration thereby making stochastic. Lillicrap et al. (2016) optimise the parameters of the Q function, 09, by optimizing for the loss
N
1
Lo
Σ
(y) - Q88 0z; 09))3
(21)
N ==1
where yi is the target value set as = ri + A0(5 ग ($1+1; 8r); 2 The algorithm updates the parameters 0 by 88 + 08 + aq oe LQ, where LQ is defined as in Equation 21 The gradient of the policy parameters is defined
1 N
V J(0))
VvR( 09) (0) (==s))
V Voo 0r) (=)),
(22)
and the parameters 0t are updated in the direction of increasing this objective.
K BACKGROUND ON SOFT ACTOR CRITIC
The goal of the SAC algorithm to train an RL agent acting in the continuous state and action but discrete time MDP M (S,A, P. Sr, 80, A), which is described in Appendix [ The SAC agent optimises for maximising the modified objective:
J(0
E.g a (s)()) + H(T(, 86; 0"))],
where H. is the entropy of the policy т. This additional entropy term improves exploration (Schul- man et al. 2017 Haarnoja al., (017). Haarnoja et al. (2018) optimise this objective by learning DNNs: the (soft) state value function V(s; 0%), two instances of the (soft) state-action value func-
tion: Q(s a,; 8 where i {1,2). and tractable policy T(s,, a,; 0*). To do they maintain a dataset 'D os state-action-reward-saalll tuples: 'D = [[ss, ai ri, s()). The soft value function is trained minimize the following squared residual error,
Jy (00) EgnD [2
Eawn [Q(s, )
log -og
(23)
where the minimum of the values from the two value functions Q; is taken to empirically estimate this expectation. The soft Q-function parameters can be trained to minimize the soft Bellman resid- ual
Jo(09) = Ex
- - -VV(::)))
(24)
where 0V are the parameters of the target value function. The policy parameters are learned by minimizing the expected KL-divergence,
((0%) BanD |KK (m(S.1;07) exp(Q(s, :00)) Zoo (8)
)]
(25)
where (s) normalizes the distribution.
L DDPG WITH GELU ACTIVATION
We provide the comparison between single hidden layer network and multiple hidden layer network because our results in Section 4 are for single hidden layer. The same architecture is used by Lill- crap et al. (2016) for the policy and value function DNNs which is two hidden layers of width 300
34

---

Published as a conference paper at ICLR 2025
(a) Walker2D
(b) Cheetah (c) Reacher
(d) Swimmer
Figure 6: Comparison of single hidden layer with GeLU activation (blue) and multiple hidden layer with ReLU activation (red) architectures for DNNs.
and 400 with ReLU activation. Here we provide the comparison to single hidden layer width 400 with GeLU activation for the architecture used by Lillicrap et al. (2016). We provide this compari- son in Figure and note that the performance remains comparable for both architectures. All results are averaged over 6 different seeds. We use PyTorch-based implementation for DDPG with modi- fications for the use of GeLU units. The base implementation of the DDPG algorithm can be found sarre:htps//ggihhub.oom http:///l/ek/////ooooooooeeeeeae as in the base implementation.
The hyperparameters are
M BACKGROUND ON SPARSE REPRESENTATION LEARNING VIA SPARSE RATE REDUCTION
We further explain and provide intuition on how the layer introduced in equation 12 learns a sparse representation. For detailed explanation, we refer to the work by Yu et al. (2023a). While equation 12 performs a linear transform to the pre-activation of the layer, it is built upon the idea of sparse rate reduction that follows the principle of parsimony: learning representations by successively compressing and sparcifying the input signal to maximally differentiate different data points for the task (Wright & Ma 2022).
first establish notation the setting of batched learning with neural networks. Suppose the objective is to represent d data points n-dimensional vectors. This problem can be formulated as obtaining a representation for d x n, denoted by which represents a batch data. The coding rate, defined by Ma al. (2007), follows,
R(Z) log det (I+ ZTZ).
(26)
To promote sparsity, which in turn learns a low-rank representation of data with a non-linear low- dimensional structure, Yu et al. (2023a) optimize the following term:
Z|o] min >| log det (I+ nee ZTzzz
max [R(Z) 2
The iterative approach learning representation that minimizes the coding rate takes gradient steps in reducing this rate. As opposed to optimizing for these objectives in a loop, as in linear programming, Yu al. (2023a) formulate this as representation learning over successive layers of NNs. Yu et al. (202aa) derive this iterative step, as representation over successive layers as follows,
4A 9a 1
z ReLUU 9(1+ DTZ
where Z are assumed to be normalised, a' is the step size, and DT is assumed to be an orthogonal dictionary for the batched dataset. A parameterised version of this transform, where both the con- stant terms interpreted being independent, is used in equation 12] The sparsity layer adds constant computation factor of 2n in the forward and backward passes, we show the impact on the steps per second metric in Figure of the appendix.
35

---

Published as a conference paper at ICLR 2025
HalfCheetah n= 2^12
HalfCheetah n 2
6000
6000
4000
2000
2000
Samples
Samples
200k
600k
800k
200k
400k
(a) log, width 12
(b) log, width 13
HalfCheetah 2-14
HalfCheetah = 2
6000
6000
4000
4000
2000
2000
Samples
Samples
200k
400k
600k
800k
200k
400k
600k
800k
(c) logg width
14
(d) log, width 15
HalfCheetah n 2^16
HalfCheetah n == 2
6000
6000
4000
4000
2000
2000
Samples
Samples
200k
600k
800k
200k
400k
600k
800k
(c) loga width = 16
(f) log, width 17
Figure 7: The canonical policy (in red) tracks the returns for linearised policy (in blue) at higher widths (log, n > 15).
36

---

Published as a conference paper at ICLR 2025
Ablation Ferupde p54
Humnanold Walk Gisusiki
6003
5033
332
3033
2003
1003
Strp
Step
(a) Ant Ablation
(b) Humanoid Stand
Humanoid
Groap Varila Graank Group Group
432
0
(c) Humanoid Walk
Figure Ablation over the aq parameter.
Dog Stand
Ant
120
120
100
100
80
60
60
40
20
20
Samples
Samples
0
100k
200k
300k
500
2M
3M
(a) Ant
Dog Stand
Figure 9: We show the steps per second for SAC (blue) and sparse SAC (red) as training progresses. We observe that, in the ant environment, for the baseline the steps per second is as high as 90 whereas in the sparse implementation this drops 65. Similarly, in the dog stand environment this difference is 81 to 61. While this increases the computation requirements and wall clock time it does not make our method intractable given the significant gains in performance.
37

---

Published as a conference paper at ICLR 2025
N FURTHER EXPERIMENTAL RESULTS
We observe that the discounted returns don't 't vary for the ant MujoCo domain Todorov et al. 2012) as shown in figure g with the environment steps on the x-axis. We see lot of variance across and within choices of aq for humanoid walk and stand environment of DM control suite (Tunyasuvu- nakool et al., 2020) even though the sparse method remains superior to SAC with fully connected feedforward. We attribute this to being an exploration problem, while our method is able to over- come couming-related bottlenecks, is unable to overcome the efficient exploration issue, which holds back the agent from attaining optimum returns in digeer-dimenioonall control tasks.
O RELATED WORK
There has been significant empirical work that assumes the set of states to be a manifold in RL. The primary approach has been study discrete state spaces data lying a graph which has an underlying manifold structure. Mahadevan & Maggioni (007) provided the first such frame- work to utilise the manifold structure of the state space in order to learn value functions. Machado iet al. (0077) and Jinnai et al. (2020) showed that PVFs can be used to implicitly define options and applied them to high dimensional discrete action MDPs (Atari games). [Wu et al. (2019) pro- vided an overview of varying geometric perspectives of the state space in RL and also show how the graph Laplacian is applied to learning in RL. Another line of work, that assumes the state space is manifold, is focused on learning manifold embeddings or mappings. Several other methods ap- ply manifold learning to learn compressed representation RL (Bush & Pineaul 2009; (Antonova iet al. 2020 Liu et al.] 2021). Jenkins & Mataric (2004) extend the popular ISOMAP framework [Tnnenbaum] 1997) to spatio-tempora data and they apply this extended framework to embed hu-
man motion data which has applications in robotic control. [Bowling et al. (2005) demonstrate the efficacy of manifold learning for dimensionality reduction for a robot's ' position vectors given addi- tional neighborhood information between data points sampled from robot trajectories. Continuous RL has been applied to continuous robotic control (Doyal 2000a Deisenroth & Rasmussen, 1 Duan et al., 2016) and portfolio selection (Wang et al. 2020; Wang & Zhou| [2020] Jia & Zhou, 2023 20226). We apply continuous state, action and time RL as a theoretical model in conjuction
with a linearised model of NNs to study the geometry of popular continuous RL problem for the
first time.
More recently, the intrinsic dimension of the data manifold and its geometry play an important role in determining the complexity of the learning problem (Shaham et al., 2015. Cloninger & Klock
2020 Goldt et al.) 2020; Paccolat et al.) 2020. [uchaaan et al., 2021 Tiwari & Konidaris] 2022)
for deep learning. Schmidt-Hieber (2019) shows that, under assumptions over the function being approximated, the statistical risk deep ReLU networks approximating a function can be bounded by exponential function of the manifold dimension. BBsri & Jacobs (2017) theoretically and empirically show that SGD can learn isometric maps from nigh-diimnssioal ambient space down to m-dimensional representation, for data lying on m-dimensional manifold, using two-hidden layer neural network with ReLUU activation where the second layer is only of width m. Similarly, Ji et al. (2022) show that the sample complexity of off-policy evaluation depends strongly on the intrinsic dimensionality of the manifold and weakly on the embedding dimension. Coupled with our result, these suggest that the complexity of RL problems and data efficiency would be influenced more by the dimensionality of the state manifold, which is upper bounded by 2d, + 1, as opposed to the ambient dimension.
We summarise several approaches for better representation learning in RL using information bot- tlenecks. Like our work, this approach reduces noise and irrelavant signal One common approach is to compress the state representation that is used by the agent for learning (Goyal et al. [219a b) 2020 Islam et al. 2022). The central idea is to extract the most informative bits with an auxiliary objective. This auxiliary objective could be exploration based (Goyal et al., 2019a)), enables hierar- chical decision making (Goyal et al. 2019b), predicting the goal (Goyal et al. 2020). and relevance to task dynamics (Islam al. 2022). While these are practical methods they do not provide the- oretical limit on the dimension of the bottleneck. In contrast, our representation is a local manifold embedding that preserves the geometry of the emergent state manifold.
38

---

Published as a conference paper at ICLR 2025
Another closely related line of research exploits the underlying structure and symmetries in MDPs. Ravindran & Barto (2001) provide a detailed and comprehensive study on on reducing the model size for MDPs by exploiting the redundancies and symmetries. There have been with other more specific approaches to this (Ravindran & Barto 2003; 2002) and more recent work follow ups by van der Pol et al. (2020). The broader study of manifolds, within differential geometry. is related to the study of symmetries and invariances. We anticipate that further reducing the effective state manifold based on redundancies, to extend our work, would be highly promising. Givan et al. (2003) and Ferns al. (2004) also provide closely related state aggregation techniques based on bisimulation metrics which have been developed further (Castro & Precup/ 2010) Gelada et al.] 2019: Zhang et al., 2020 Lan et al. 2021). The bi-simulation literature defines metrics that incorporates transition probabilities or environment dynamies of the environments. The underlying metric probabilistic in nature. The manifold and metric are defined in such way as to facilitate better representation learning for RL. The primary difference is that our approach proves how a low-dimensional manifold "emerges" from the design and structure of certain continuous RL problems.
We finally contextualize our work in light of various control theoretic frameworks. Control systems on non-linear manifolds have been studied widely (Sussmann, 1973; Brockett 1973, Nijmeijer & van der Schaft [990] Agrachev & Sachkov [2004] [Bloch & Bloch| 2015 Bullo & Lewis] (019))
Like most control theoretic frameworks the transitions, dynamies, and the geometry of the system are assumed known to the engineer. (Liu et al.] (021) recently provide a framework for controlling a robot on the constraint manifold using RL, As noted previously, our work is also closely related to the notion reachability in control theory (Kalman 1960; Jurdjevic 1997, Touchette & Lloyd 1999, 2001) which deals with sets reachable under fixed and known dynamics of system. Reachability sets from control theory have been applied for safe control under the RL framework (Akametalu jet al., 2014; [Shao et al. 2020; Isele et al.] 2018). While the objective is similar, to find the sets of states reached with any controller, the assumptions, on the underlying dynamics are different leading to different results.
P EXTENSION TO OTHER ACTIVATIONS AND ARCHITECTURES
It is difficult to theoretically analyse complex engineered systems such as neural networks for continuous control learned using policy gradient methods. We have simplified this setting by us- ing linearised policies (Section] 3.I) with GeLu activation and access to the true gradient of the value function (Section 2.1. We show results in GeLu activation because it is the closest (smooth) analogue to the most popularly used ReLuu activation which very commonly used in continuous control with RL (Lillicrap) 2015) Schulman et al., 2017] Haarnoja et al. 2018). Despite our choice of GeLu, as noted in Section 4 our results extend to activations which are twice differentiable ev- erywhere with bounded derivatives. Moreover, our results capture the setting of neural policies that have a very high-dimensional parameter space but also have structured outputs (Lee et al. 2017 Ben Arous et al. 2022). In study of supervised deep learning results emanating from theoretical models that approximate shallow wide NNs have been extended to deeper networks, c.g. the neural tangeent kernel (NTK) framework (Jacot et al., 2018). Moreover, there have also been mechanisms to make finite depth and width corrections to NTK (Hanin & Nica 2019). Theoretical inferences made in simplified settings have been extended to applications and wide range of architectures as well (Yang& & Hu 2021; Yang et al. 2022, Fort et al. 2020 Wang et al. 2022). We anticipate that extending our results to a broader set of activations, architectures, and reinforcement learning algorithms would lead to better applications by means of improved theoretical understanding.
Another assumption we make deterministic transitions. While this is true in many popular bench- mark environments (Todorov et al. 2012 Tunyasuvunakool et al. 2020). the most general setting of RL as model for intelligent agent the transitions are stochastic. This is a common feature in control theory where results in deterministic control: s(t) = g(s(t), u(t), t), with continuous states, actions, and time, can be extended to stochastic transitions by considering bounded stochastic perturbations
s s(t) = g() u(t), t) + d(s(t), u(t), t)dw,,
where d is the stochastic perturbation aspect with w, being the Wiener process and u(t) is the open loop control. Tor example, the contraction analysis by Lohmiller & Slotine (1998) in deterministic transitions is extended to stochastic perturbations by Pham et al. (2009). We anticipate that our analysis, under appropriate assumptions on the stochastie perturbations, has promise of extensions.
39


Reviews:
[{"review_id": "z92YsxcLP9", "reviewer": "Reviewer_aRZw", "summary": "This paper studies the Manifold hypothesis for deterministic continuous state and action environments in RL. In particular they show that the set of reachable states for a certain class of analytically tractable neural network policies learned using stochastic policy gradient lies on a manifold with a dimension upper bounded by a linear function of the action space dimension. They demonstrate this upper bound in empirical experiments and show that using a sparse representation can help in learning in complex environments.", "questions": "- In Figure 2, I thought that the upper bound is supposed to be $2d_a + 1$ but the caption says that the green line is $d_s$ and the estimated dimensionality is below it. Also what is the red line in this figure?\n\n- In Figure 4, what is “Group” in the legend referring to?", "limitations": "", "rating": 8, "confidence": 2, "soundness": 4, "presentation": 3, "contribution": 4, "strengths": "The paper is well-written and rigorous. \nThe analysis of the manifold hypothesis for RL using neural network policies is novel to my knowledge, and a significant contribution to the area.", "weaknesses": "I don’t have any major comments on the weaknesses of the paper, I feel that the authors adequately mention the limitations in Section 6.  \n\nMinor:\n- There were a few awkward or incomplete sentences that I did not understand:\n  - L263-265: “In theoretical frameworks ….”\n  - Caption of Figure 3\n  - L462: “A common of a fully …”\n  - The paragraph under Equation (14), and specifically the sentence in L475-478.\n- The fonts in Figures 2,3, and 4 were too small to read.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "M7QCV1gt5A", "reviewer": "Reviewer_zo6o", "summary": "The paper presents a novel theoretical and empirical analysis of reinforcement learning (RL) within continuous state and action spaces by employing a geometric framework. The authors propose that the set of attainable states in RL systems is constrained within a low-dimensional manifold, which is directly influenced by the dimensionality of the action space. Specifically, they examine the structure of locally attainable states using semi-gradient updates in actor-critic learning, aiming to identify and exploit low-dimensional representations for RL in high-dimensional environments. This approach is validated through experiments on MuJoCo environments, demonstrating the practical implications and performance gains of using low-dimensional manifold learning in RL settings with high degrees of freedom.", "questions": "- Q1: Is the sparsification layer computationally efficient compared to fully connected layers, particularly as the number of states and actions grows?\n\n- Q2: How does this work compare in performance with other state representation techniques, such as latent variable models in RL?\n\nIn overall, the paper is generally well-written, with each section following logically from the previous one. However, the introduction to differential geometry concepts (e.g., Lie derivatives, exponential maps) could be more concise, as this level of detail might obscure the main contributions. The plots provided for MuJoCo environments are helpful, but additional figures explaining the theoretical framework would improve clarity.\n\n\nThis submission is a strong contribution to theoretical RL research and provides significant insights into the geometric structure of attainable states in continuous environments. The empirical validation in MuJoCo environments further solidifies the practical value of the proposed method. \n\n* Minor comments\n\n- Figure 2: What is the red line?", "limitations": "", "rating": 8, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- Strong theoretical insights: The authors look at the geometry of RL dynamics in continuous-time MDP with continuous spaces in order to link the dimensionality of the attainable state manifold to the action space. Theorem 1 is the main contribution in this reward. The theorem formally shows that the dimension of this manifold is related to the dimensionality of the agent's action space rather than the full state space. Specifically, the dimension of the manifold is approximately $2 \\times (action dimension) + 1$. I found this insight quite interesting. It looks like the theorem suggests that by understanding this low-dimensional structure, we can design more efficient policies or networks, as we don’t need to account for the full high-dimensional state space, as seen in results in section 5.3.\n\n- Interesting empirical experiments on multiple MuJoCo environments that support the theoretical claims (see above comment). By measuring the dimensionality of the attainable state space, the results indicate consistency with the derived bounds, highlighting the empirical validity of the proposed approach.", "weaknesses": "- The mathematical presentation is too complex: The derivation and presentation of the theoretical framework are mathematically dense, which could limit accessibility for a broader audience. The notation and terminology, especially in the sections on Lie series and vector fields, may be difficult for readers less familiar with differential geometry. It would be great if there could be some better insights following each main step.\n\n-  There is little discussion on scalability: While the results are validated empirically, the paper could benefit from further discussion on the computational complexity and scalability of the approach, particularly in scenarios with much larger state-action spaces or image input space.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "WGL2OFyb9H", "reviewer": "Reviewer_S4jN", "summary": "The paper proposes a theoretical approach to understanding reinforcement learning in continuous state and action environments using a geometric perspective. Unlike traditional theoretical RL models that focus on finite state and action spaces, this work introduces the concept of a low-dimensional manifold that captures the space of states achievable by RL agents trained with actor-critic methods. By demonstrating that the dimensionality of this manifold is bounded by the dimensionality of the action space, the authors establish that RL agents operate within a constrained subset of the full state space.", "questions": "How could the theoretical framework be adapted to account for stochastic transitions and noisy reward functions, which are common in real-world environments?\nCan the manifold-learning approach be tested with deeper and more complex neural network architectures to better reflect the structure of contemporary RL models?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "By proving that RL agents operate within a low-dimensional manifold of attainable states, the authors offer a mathematically rigorous insight that connects the geometry of state spaces to the action dimensionality. \nThis theoretical framework is supported by empirical evidence from simulated environments, such as MuJoCo, showing that training dynamics in RL indeed produce a low-dimensional representation. The paper presents a practical application by incorporating a manifold-learning layer in policy and value networks, which improves performance in complex control tasks.", "weaknesses": "The analysis assumes deterministic transitions and access to an exact value function, which is often impractical in dynamic, stochastic environments. \nThe simulation setups may not fully capture the complexity and variability of real-world tasks where environmental noise and high-dimensional data structures can complicate learning dynamics. \nLastly, the mathematical framework is limited to two-layer neural networks, which may oversimplify the behaviors of deeper architectures commonly used in modern RL, potentially limiting the generalizability of the findings to more complex neural network models.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "CabqxPBCOD", "reviewer": "Reviewer_7BEr", "summary": "This work shows that training dynamics of a two layer neural policy induce a low dimensional manifold of attainable states embedded in the high-dimensional state space.", "questions": "1. What does $\\mathbb{B}_{W_{k \\eta}}$ below equation (9) represent? Why does batch data need to come from SDE? Which research works on continuous time policy gradients have used this concept? \n2. Does $G(W)$ in equation (10) contain $\\eta$? \n3. Since similar proof techniques were used, what are the innovative points of this manuscript compared to Ben Arous et al. (2022)?\n4. The paper studied MDP for continuous-time systems. However, the simulation case is MuJoCo for discrete-time systems. \n5. The gradient of policy network parameters for DDPG and SAC algorithms in simulation work deviates from the continuous time policy gradient. To what extent can simulation work verify the effectiveness of theoretical results?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 3, "strengths": "theoretical understanding", "weaknesses": "1. The manuscript is difficult to follow. Section 2 introduces a large number of mathematical concepts, including manifolds, fields, and Lie algebras. However, the connections between these mathematical concepts are fragile. I did not see the main storyline and the author's contribution clearly before page 7. \n2. The paper primarily considers two-layer neural networks with GeLU activation. However, the subsequent theoretical analysis depends on linear approximation. I cannot tell the difference between this and single-layer neural networks and the impact of approximation errors on the proof results. Besides, extending the analysis to other network architectures and activation functions would strengthen the claims of universality.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Oral)