Geometry of Neural Reinforcement Learning in Continuous State and Action Spaces

Thank you for your time studying our manuscript and giving your valuable feedback. We address each concern individually below.

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.

This issue was also pointed out by other reviewers. We have whittled down the contents of section 2 to focus on the relevant details by moving some text to the appendix. Moreover, we have added explanations on how each one of these concepts relate to RL. Our main result section now starts on page 6 and we also provide context on how these results relate to the mathematical concepts explained in the background sections. We also note that despite our main result in section 4 we also introduce our learning setup in section 3. Please let us know if it addresses your concerns.

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

As pointed out by you, we show results in GeLu activation because it is the closest (smooth) analogue to the most popularly used ReLu activation. The linearised model we use is applied very commonly as a tractable approximation of neural networks (NNs) [1,2,3]. We refer to two layer neural networks which are the same as single hidden layer neural networks. Is that what you mean by a single layer NN? These are different from single layer NNs because the output is linear in NN parameters but not in the input, whereas a single layer linear NN is linear in the input. Learning in the wide linearised model of a two layer NN can be interpreted as the agent “picking and choosing” from a large number of randomly projected “features” to fit to the task, which in this case is an optimal policy. Moreover, Jaehoon et al [1] show that for supervised learning in the limit of width going to $\infty$ both the outputs and loss converge for training (Theorem H.1 and Section 2.4 “Infinite width networks are linearized networks” from their paper) and they back it up empirically (Figures S8, S9).

We have now added an empirical comparison of discounted returns, for the Cheetah environment learned using a DDPG agent, between the linearised and canonical two layer NNs for increasing width in Section 5 (Figure 2) and report discounted returns as width increases in Figure 7 in the appendix . We observe that both the linearised and canonical neural networks have the same returns, across episodes, at very large width. Intuitively, this means that a linearised NN behaves similar to a canonical two layer NN when optimizing returns using a policy gradient algorithm. We are working on a theoretical result on how this approximation impacts our proof, on the estimated dimensionality, and we will post it as soon as we have it. Please let us know if this sufficiently addresses your concerns?

Besides, extending the analysis to other network architectures and activation functions would strengthen the claims of universality

We note in section 4, that we require the following restrictions on the activation function:

Furthermore, we assume that the activation, $\varphi$, has bounded first and second derivatives everywhere in $\mathbb R$.

Our results apply to any activation which satisfies this condition (e.g. tanh). We have now added this explainer in the Appendix under the section: Extension to Other Architectures and Inputs. We also note how our results can lead to similar results for deeper fully connected networks. Our current work opens up room for obtaining theoretical results that apply to new architectures. Moreover, deep learning theory results in wide fully connected neural networks have been applied to obtain better theoretical understanding and applications of other neural network architectures [4,5,6,7,8]. We hope our results would further similar understanding of deep reinforcement learning which we now explain and note in the aforementioned Appendix section. As you might have noticed, our current manuscript is sufficiently explanatory of empirical settings and complex as is and extending it to more architectures might seem like a bit of an over stretch.

Reviewer 7BEr: if possible, can you respond to the rebuttal?

Thank you for carefully going through our work and your supportive comments. We have tried to address in a lot more detail the primary weakness pointed out by:

The simulation setups may not fully capture the complexity and variability of real-world tasks …

We have added a section in the appendix titled “Extension to Other Activations and Architectures” on how applicable our results are to more complex neural network architectures and stochastic transitions. In summary, we remark how in the past theoretical analysis of shallow large width networks [1] have been extended to finite depths and widths [2]. We also remark how continuous control results from deterministic dynamics [3] have been extended to stochastic dynamics [4]. We make these arguments in much more detail with various examples.

We recognise that this is a valid limitation of our work but there is precedent that simplified theoretical analyses leads to broader application.

We answer your questions below.

How could the theoretical framework be adapted to account for stochastic transitions and noisy reward functions, which are common in real-world environments?

Many popular RL benchmarks for continuous control have deterministic transitions (e.g. MuJoCo and DM Control). We agree that the more real world cases will be stochastic. Presenting theoretical results with deterministic transitions and extending them to stochastic transitions is more common. One such model is to provide results for the deterministic case $\dot{s}(t) = g(s(t), u(t), t)$ and then extend to the case with stochastic perturbation: $\dot{s}(t) = g(s(t), u(t), t) +d(s(t), u(t), t) dw_t,$ where $d$ is the stochastic perturbation aspect with $w_t$ being the Wiener process and $u(t)$ is the open loop control. We now make this remark in the appendix.

Can the manifold-learning approach be tested with deeper and more complex neural network architectures to better reflect the structure of contemporary RL models?

Many popular continuous control RL frameworks with sensory inputs utilise three layer NNs, which is what we benchmark against in action 5.3. At the same time contemporary methods employ latent transition models (e.g. Dreamer v3 [5]) which are learned using complex deep NNs. One possible application of our work could be to incorporate the low-dimensional aspects of the attainable set of states into these transition models. We also anticipate significant performance improvement in future work with deeper networks constructed using the sparsification layer we have used.

References:

[1] Neural Tangent Kernel: Convergence and Generalization in Neural Networks, Arthur Jacot, Franck Gabriel, Clément Hongler

[2] Finite Depth and Width Corrections to the Neural Tangent Kernel, Boris Hanin, Mihai Nica

[3] On contraction analysis for non-linear systems. Winfried Lohmiller and Jean-Jacques E. Slotine

[4] A contraction theory approach to stochastic incremental stability. Quang-Cuong Pham, Nicolas Tabareau, and Jean-Jacques Slotine

[5] Mastering diverse domains through world models, Hafner, Danijar and Pasukonis, Jurgis and Ba, Jimmy and Lillicrap, Timothy

Reviewer S4jN: if possible, can you respond to the rebuttal?

Thank you very much for your support and approval of our work. We have fixed the typos and mistakes you have mentioned. Once again, your valuable comments are much appreciated. We have also fixed the paragraph in section 5.2 as below:

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 $t$ under the dynamics $\dot{s(t)} = A s(t) + B \Phi(x) W$, where $A, B$ are as in equation above. To achieve this, for fixed embedding dimension $d_s$ we obtain neural networks sampled uniformly randomly from the family of linearised neural networks as in definition above, with $r = 1.0, t \in (0, 5), n = 1024$. Consequently we obtain 1000 policies with $\delta t = 0.01$, and therefore a sample of 500000 states to estimate the intrinsic dimension of the attained set of states using the algorithm by Facco et al, 2017. We vary the dimensionality of the state space, $d_s$, from 3 to 10 to observe how the intrinsic dimension of the attained set of states varies with the embedding dimension while keeping $d_a$ fixed at 1. The dimensionality of the attained set of states remains upper-bounded by $2d_a + 1 = 3$ for this system (figure 3).

We have increased the font size on the images as well. Answer to your questions:

Q1: We have fixed that error, the red line is now at $2 d_a + 1$ it was previously at $d_a + 1$ which was an oversight on our side. The green line still remains at $d_s$ to show the discrepancy between the state space dimensionality and the attainable set dimensionality.

Q2: Group in the legend is residue from the wandb interface which we use for tracking our learning algorithms; we have now removed it. Thank you for pointing it out.

Reviewer aRZw: if possible, can you reply to the rebuttal?

Thank you for your valuable and insightful review and pointers on how to improve our work. We have significantly re-organised the background section and also provide better insights following our main theoretical result. More specifically, we intuit what the bases of this low-dimensional manifold are. We have also added a discussion on scalability and address it in detail below.

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.

We apologise for the oversight in this aspect. We have now moved text that wasn't absolutely needed to understand our results and the main story to the appendix. We have also additionally provided intuitive explanations in sections 2.2 and 2.3 and how they relate to RL. We have also added a broader more intuitive proof sketch. Please take a look and let us know if this addresses your concerns.

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.

The sparsity layer (equation in section 5.4) adds a constant computation factor of $2n^2$ in the forward and backward passes where $n$ is the width of the layer, the big-O computational complexity remains the same though. This is reflected in the metric steps per second. In the most recent version we report this in the Appendix section titled “Further Experimental Results”. 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 to ~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. We expect this metric to be impacted as the input size increases.

Image inputs are significantly different, in our opinion. The theoretical setting is that of the input size going to $\infty$ with the width (see chapter 2 in textbook by Romain and Zhenyu [1] or [2,3]). This is so because in the very high dimensional setting the input size (e.g. 256 X 256) is comparable to large widths of neural networks. We assume the input size is orders of magnitude smaller, which is true of control/sensor inputs in MujoCo and deepmind control suite.We make a note of this distinction in our discussion section but as of now we do not delve into how the manifold dimensionality and subsequent representation learning will change in this setting.

We have answered Q1 above. For Q2:

Q2: How does this work compare in performance with other state representation techniques, such as latent variable models in RL?

Do you mean latent state models like dreamerV3 [5] and such? Dreamer does not perform well, in fact learning does not progress at all, for dog and humanoid environments of deep mind control suite (see Figure 5 in [4]). These are particularly “notorious” high-dimensional environments where we show great improvement. We also note that model based methods with significant added complexity do perform better whereas we have simply made the neural net wider and modified a layer while keeping the number of parameters the same as in the baseline. Please let us know if that is not what you meant by “latent space models” or if we have not addressed your concern.

Figure 2: What is the red line?

This was an error, thank you for pointing it out. We drew a red line at $d_a + 1$ but it was supposed to be at $2 d_a + 1$. We have made that fix with an explanation in the latest version of our paper.

References:

[1] Random matrix methods for machine learning, Couillet, Romain and Liao, Zhenyu

[2] Phase diagram of stochastic gradient descent in high-dimensional two-layer neural networks, Veiga, Rodrigo and Stephan, Ludovic and Loureiro, Bruno and Krzakala, Florent and Zdeborov{'a}, Lenka

[3] Dynamics of stochastic gradient descent for two-layer neural networks in the teacher-student, Goldt, Sebastian and Advani, Madhu and Saxe, Andrew M and Krzakala, Florent and Zdeborov{'a}, Lenka

[4] Td-mpc2: Scalable, robust world models for continuous control, Hansen, Nicklas and Su, Hao and Wang, Xiaolong

[5] Mastering diverse domains through world models, Hafner, Danijar and Pasukonis, Jurgis and Ba, Jimmy and Lillicrap, Timothy

Reviewer zo6o: if possible, can you reply to the rebuttal?