Wide Neural Network Training Dynamics for Reinforcement Learning

While I found the paper to be an interesting read I observe some crucial issues with the work which need addressing.

Lack of related work: While the authors have cited relevant work in DL theory, they are missing come crucial citations [1, 2, 3, 4]. The issue is that previous works have explored similar problems. For example, the setting where gradient "flows" through both $Q_{\theta}(s, a)$ and $Q_{\theta}(s', a')$ is called residual learning [3, 4]. The equation number 7 by the authors has a single sample analog in [3] (see Equation (1) in the paper). Baird [4] introduced the idea of residual algorithms to make the TD algorithm more principled. Further see related work [5]. Please let me know if I am mistaken about this connection. [1] has a different type of linearization: where they linearize the learnt function in the first layer weights for single hidden layer networks, while this is not directly applicable it has a connection with the current work. Similarly, [3] gives insights into when and if TD with NNs and homogenous activations converges albeit in discrete setting, depending on the environment.

Not a substantial enough conclusion: I appreciate that the authors have derived and written down the dynamics of learning in the NTK regime. While it is a useful contribution, I am afraid that I am not able to discern a valid and strong conclusion.

I see that the authors have introduced the TD version of $\Theta_0$ (just like Lee et al. 2019) but I would be interested in how this matrix effects the convergence. I believe this is where the environment dynamics come into play. Further, I would expect a theorem similar to Theorem 2.1 by Lee et al. 2019 (latest arxiv version). Theorem 2.1 by Lee et al. 2019 provides a sense of how far the linear approximation $f^{lin}$ is from the true function $f$ i.e. order of $1/n^{-1/2}$. I would expect something similar with relation to the environment dynamics (and not only the width of the NN) for the work to be more applicable to TD learning.

Similarly, Brandfonbrener and Bruna, 2020 [2] provide some insights into when and if neural TD learning converges depending on the geometry of the space of function approximators, the structure of the underlying Markov chain, and their interaction. This brings in aspects of reinforcement learning and transition dynamics. This issue also relates to the previously stated issue of a lack of related work.

Further, the experiments seem a bit inconclusive to me, especially the experiment on cartpole domain. While I appreciate the new comparisons in Figures 3 and 4, even in a toy environment we see the parameters behaving differently (in Figure 1). The authors also remark that their analytic model does not match the empirical model in Section 5.1 but no reasoning is provided. For cartpole the result is surprising the linearized dynamics work better and is definitely a meaningful contribution. I would be curious to see results normalized by the true values? As far as I can tell, both NNs and TD-NTK do not "fit very well" with different degrees of "not fitting very well".

Statement of assumptions: I would be curious to know when TD-NTK does not converge? I would be surprised if it converges everywhere without assumptions (as I presume is noted in Equation 31 in the appendix). For example, [5] and [4] both say that residual RL might not always converge. Similarly, Lee et al 2019 assumed $\lambda_{min}(\Theta_0) > 0$. Further discussions regarding assumptions can be found in [1] and [2]. These assumptions are better stated clearly.

I do not intend to discourage the line of work the authors are pursuing. It is very important for the research community that we have a concrete theory of Deep RL that relates to practice. The training dynamics derivation is very useful and much appreciated but the work would fair better with a stronger conclusion and connections to previous work.

---

References

[1] Neural Temporal-Difference and Q-Learning Provably Converge to Global Optima, Qi Cai, Zhuoran Yang, Jason D. Lee, Zhaoran Wang, 2019

[2] Geometric Insights into the Convergence of Nonlinear TD Learning, David Brandfonbrener, Joan Bruna, ICLR 2020

[3] Deep Residual Reinforcement Learning, Shangtong Zhang and Wendelin Bohmer and Shimon Whiteson, 2019

[4] Residual algorithms: Reinforcement learning with function approximation, ICML 1995

[5] Gradient temporal-difference learning algorithms, Richard S. Sutton and Hamid Reza Maei, 2011

Although the studied research direction is interesting and promising, the paper has some significant flaws, and the contribution of the paper is weak in its present form. The study requires a significant extension to be further considered for publication.

1. The theoretical contribution of the paper is merely the calculation of NTK dynamics for Q-learning by supervised learning over fixed datasets of rollouts adapted from existing work [ Lee et al. (2019) ]. This part does not bring any new contributions or insights.

2. The title, abstract, and introduction of the paper need to be more accurate. Setting studied in the paper are not those of online or offline reinforcement learning. The studied setup can be characterized as a Q-learning approach via supervised learning over a fixed dataset. Namely, a dataset of dumped policy-environment interactions is used to train the parameters of a Q-network. It differs from the standard online and offline actor-critic RL setups in which the actor (policy) and critic (Q-network) is approximated simultaneously. Namely, the work ignores the issue of approximating the policy action on the next-state; this kind of analysis detached from policy properties is a pure Q-learning. As stated, the Bellman bootstrap equation reads $Q^{\pi}(s, a) = r + \gamma \sum_{s'} P(s'|s,a) \pi(a'|s') Q^{\pi}(s', a')$, but it does not lead to the supervised learning loss function (6). $\mathcal{X}'$ is unavailable in the RL setup. The dataset MDP tuples $(s_i, a_i, r_i, s_i')$ does not contain the next actions; they should be obtained from the policy instead, as the Q is an estimated value of the trajectories obtained by following the policy $\pi$. It is a separate issue of how to regularize and avoid issues with distribution shift at this stage addressed by modern Offline RL methods, including TD3-BC [1] and FDCR [2].

[1] Kostrikov I., "Offline Reinforcement Learning with Fisher Divergence Critic Regularization", ICML 2021 [2] S. Fujimoto and S. Gu, "A Minimalist Approach to Offline Reinforcement Learning", NeurIPS 2021

3. Experimental evaluation could be stronger. Only simple environments (a toy one and a cartpole are used) are particular, and it is unclear how the obtained conclusions will transfer to more complex environments widely used as benchmarks (MuJoCo).

Minor remarks:

• p.5 below (11) "maximum and minimum values of Θ." -> eigenvalues,
• p.5 below (11) "This is expected because in this regime learning the value function is essentially a supervised learning task of learning the immediate rewards, and all terms with policy dependence (those involving X ′ ) drop out as would be expected in a supervised learning task." in supervised learning task that is true, but unfortunately not in an RL task :)
• p. 6 (5.1) "the value will be $1$ (taking $\gamma = 1$) for all $s < 1$ and $0$ otherwise." I think the condition is switched here (value $0$ for $s<1$).
• p. 8 (5.2) "the true returns obtained by rolling out the evaluated policy in the environment to evaluate our method." the issue of policy-generating data requires more attention in the experimental section.
• derivations in C.1 and C.2 are not needed, these are solutions of simple non-homogenous differential equations,

While in general the paper presents useful findings, I think it may lack enough overall contribution for publication at ICLR. It's not clear what the key takeaways are from the paper beyond the differential equation solutions, which are closely based on prior work. Some specific reasons that I feel there is not enough contribution:

• The theoretical results are presented without much explanation or intuition. It's not clear if there are any insights that can be gained from the solutions to the differential equations.
• The empirical results are only for the simplest of RL environments (1 and 4 dimensional state spaces; 0 and 1 dimensional action spaces). Since deep RL is primarily useful for high-dimensional environments with complex value functions, it would be more convincing if the authors evaluated on these types of environments as well.
• The analysis is restricted to finding the Q-function of a fixed policy, while reinforcement learning algorithms generally use a TD loss like $\left(Q(s, a) - (R(s, a) + \\max_{a'} Q(s', a'))\right)^2$.

Also, there has been some previous analysis of RL with NTK approximation by Yang et al. [1], who go beyond training dynamics for a fixed-policy and give regret bounds for an optimistic variant of Q-learning with the NTK. It would be good to include a comparison to their work.

[1] Yang et al. On Function Approximation in Reinforcement Learning: Optimism in the Face of Large State Spaces. NeurIPS 2020.

The paper mainly focuses on the training dynamics of full-gradient TD-learning, which, to my understanding, is similar to minimising the average TD-error. This severely limits the potential impact of the work since full-gradient algorithms are not used in practice and, more importantly, it's not entirely clear if the solutions to the objectives in full-gradient algos are actually desirable (see eg. chapters 11.4 - 11.5 in Sutton). I believe this aspect should be better highlighted and contextualized in the paper.

Authors also show that their NTK derivation of TD-learning with semi-gradient descent and target network (the more relevant setup) does not follow the same dynamics as under SGD which is an interesting finding even if negative. However the relevant figures are pushed to the appendix and there is not discussion as to why might this be the case. Is it because NTK simply is not a good approximation for algorithms which do not follow a proper gradient field? Is it because of the added complexity of the slowly moving target network?

The last question also points to some lack of clarity in the paper as it bundles together target networks with semi-gradient methods. However having a different target function is a separate design decision from choosing not to use the gradient of the target network. This also leads to a possible experiment: is the linearisation better at approximating training dynamics in the semi-gradient setup without a separate target network?

Finally, it would be interesting to mention and develop how your empirical findings impact some of the previous works that use the linearisation of neural (gradient/semi-gradient) TD-learning to come-up with convergence proofs. Examples include:

• https://proceedings.neurips.cc/paper/2020/file/75ebb02f92fc30a8040bbd625af999f1-Paper.pdf
• https://arxiv.org/pdf/1912.04511.pdf