Paper: Wide Neural Network Training Dynamics for Reinforcement Learning (ICLR-2024)

Reviewer comment (Theory):
Lack of analysis on how the TD version of $\Theta_0$ affects convergence: The authors have introduced the TD version of $\Theta_0$ (just like Lee et al. 2019) but have not analyzed how this matrix affects convergence. The reviewer believes that environment dynamics are important here and expects a theorem similar to Theorem 2.1 by Lee et al. 2019, relating the linear approximation $f^{lin}$ to the true function $f$, but in a way that accounts for environment dynamics as well as NN width.

Author response:
We have not included a thorough derivation of the approximation $f^\text{lin}$ distance to $f$ similar to theorem 2.1 by Lee et al. (2019); as we describe in the text (section 4.1, after eq. 6) why the reasoning is analogous if the equivalent assumptions from Lee et al. are made in the temporal difference setup, and the same results, such as errors of the order $1/n^{-1/2}$ with hidden layer with $n$, still holds.

We do agree with the reviewer that the environment dynamics manifest themselves in $\Theta_0$ which, as we point out in eq. 11, would affect what the critical learning rate for convergence is. As far as training dynamics are involved, we don’t think the environment dynamics (at least in the deterministic dynamics case) introduce additional caveats towards the conditions necessary for our approximations to hold.

Reviewer comment (Evaluation):
Inconclusive experiments: The experiments, especially on the cartpole domain, seem inconclusive. The reviewer appreciates the comparisons in Figures 3 and 4, but notes that even in a toy environment, the parameters behave differently (Figure 1). The authors remark that their analytic model does not match the empirical model (Section 5.1) but provide no reasoning for this. For cartpole, the result that linearized dynamics work better is surprising and meaningful, but the reviewer would like to see results normalized by the true values, as both NNs and TD-NTK do not seem to fit very well, with varying degrees of mismatch.

Author response:
Answers to Questions:
The cartpole environment has a 4-dimensional state space (position of the cart, velocity of the cart, angle of the pole, angular velocity of the pole). We use 2048 transitions $(s, a, r, s’)$ to train the networks.
The behavior policy was a $\max_{a’} Q(s’, a’)$ from a Q-function learned by a single hidden-layer DQN agent trained online.
Finding an expression at arbitrary state-actions $\mathcal{X}^*$ is indeed to investigate the generalization of the predictions of the model/ensemble. The high-level idea to derive this is that by the linearisation of $Q(\mathcal{X^*}) = Q_0(\mathcal{X}^*) + \nabla_\theta Q_0(\mathcal{X}^*)\omega_t$, one can substitute the expression for $\omega_t$ to obtain the expression for Q-values at arbitrary state-actions $\mathcal{X}^*$ throughout and at the end of training.
That is true, some degree of discrepancy is expected to accumulate throughout training.
The environment represented in figure 2 is the toy environment, having only one action, although we can include this for the cartpole environment

Reviewer comment (Novelty):
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.

Author response:
We want to first maintain that our work is distinct from supervised learning. We believe it is inaccurate to frame our approach as “supervised Q-learning” and disagree with the reviewer in our theoretical analysis not providing novel contributions. Supervised learning assumes ground-truth labels available for the quantity that is being learned, and we stress is not the case in our setup: we do not assume no return samples or extended traces are available (only states, actions, next-states and immediate rewards for single transitions are observed). The novelty in our contribution is exactly to describe how to adapt the methods that analyze training dynamics from the neural tangent kernel perspective in supervised learning to the temporal difference policy evaluation setting with only $(s,a,r,s’)$ tuples available. Crucially, the policy being evaluated does not need to coincide with the policy that generated the data.

Reviewer comment (Writing):
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.

Author response:
Again, we disagree that the setting considered is supervised learning as no ground-truth return labels are available for training. We agree with the reviewer that we are not directly modeling the full “reinforcement learning” task, and this is highlighted in the abstract where we specifically describe our contributions as being relevant to “temporal difference policy evaluation methods” and associated uncertainty quantification. Nonetheless, forms of temporal difference off-policy evaluation are key elements present in many RL algorithms, and we believe better understanding temporal difference policy evaluation with neural networks is an important milestone towards better understanding deep RL (particularly in the offline setting with a static dataset).

We would also like to point out that, as specified in section 3, the actions in the state-action pairs $\mathcal{X’}$ are those computed from the policy *that we wish to evaluate* at the next-states, which is assumed given in the policy evaluation setting and would correspond to the action taken by the actor in the RL setting (in the RL setting, this quantity is also derived from the actor). In fact, we entirely agree with the reviewer’s comment that “Q is an estimated value of the trajectories obtained *by following the policy* $\pi$”. As the action in $\mathcal{X’}$ is taken from the policy we are evaluating, and is not related to the next actions taken in the trace that generates the data, these are exactly the Q-values learned by training with the loss in eq. 6.

Reviewer comment (Theory):
On p.5 below (11), the statement "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 tasks that is true, but unfortunately not in an RL task.

Author response:
We would like to note, that for the setting where $\gamma = 0$ (minor remark #2), the equation simplifies to the supervised learning case. The other case, when $\gamma \neq 0$ it is no longer a supervised learning task, and we wanted to highlight this difference in the text.

Reviewer comment (Theory):
On p.6 (5.1) "the value will be \(1\) (taking \(\gamma = 1\)) for all \(s < 1\) and \(0\) otherwise." — the condition appears switched here (value \(0\) for \(s < 1\)).

Author response:
For minor remark #3 we maintain that the true value function will be 1 when $s$ is less than 1 ($V(s) = 1$ when $s < 1$). The agent always moves right, and crossing the line at $s=1$ grants a reward of 1, so everything on the left of the line ($s<1$) has value 1 (when $\gamma=1$) and everything to the right ($s>1$) leads to terminal states with no further rewards, and the value here must be 0.

Reviewer comment (Theory):
The theoretical results are presented without much explanation or intuition, and it is not clear if there are any insights that can be gained from the solutions to the differential equations.

Author response:
We are grateful to the reviewer for their time and effort in providing constructive feedback to our work. We also appreciate the author acknowledging the novelty in our analysis and the potential it brings for insight into deep RL theory.

**Question 1 (What insights can be drawn from the analytical solutions for neural tangent kernel TD learning gradient flow?):**

The regime (i.e., when learning rate $\eta < \eta_\text{critical}$) suggests that many NNs used in practice may be operating in the setting we are investigating. Also, the derived kernel $\Theta$ captures the training dynamics of these settings, and can be studied in the future for insight into stability and generalization. We think that analyzing the behavior of infinitely-wide NNs offers a foundation for developing a systematic comprehension of finite-width NNs.

Reviewer comment (Experiments):
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.

Author response:
**Question 2 (Are the analytical solutions presented in this paper also good approximations for TD learning in higher-dimensional state/action spaces?):**

The derived solutions in theory are general, and also apply to higher-dimensional state/action spaces. While we appreciate that this is not confirmed with empirical experiments in the current version of our work, our focus was on exposing the theory, and much like previous work investigating neural tangent kernel training dynamics we do not expect size of state or action space size to have a significant effect on the quality of the training dynamics approximations

Reviewer comment (Theory):
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$.

Author response:
**Question 3 (Can the analysis be extended to using a fitted Q-iteration loss where $Q(s’, a’)$ is replaced with $\max_{a’} Q(s’, a’)$?):**

We appreciate that this is a very relevant extension to our work which we are actively investigating, and that it may introduce additional insight into the instabilities present in offline RL related overestimation of Q-values.

Reviewer comment (Novelty):
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.

Author response:
**Concern 1 (are the solutions to the objectives in full-gradient algos actually desirable?):**

While we agree that full-gradient algorithms are not yet often used in practice, it is a common simplifying assumption in analyzing training dynamics (for example in the neural tangent kernel literature) to retain a degree of theoretical tractability while still maintaining the core ingredients of the optimization process involved. While we appreciate the example highlighted by the reviewer as to why the stated gradient objective may not always be desirable, these are valid for some stochastic MDPs, and in many other cases, including all environments with deterministic transitions, the training will converge to a consistent value function. We have amended the manuscript to discuss these two points, and therefore explain why our approach is still sound for many environments of interest.

Reviewer comment (Presentation):
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 no discussion as to why this might be the case.

Author response:
**Concern 2 (no discussion on why TD semi-gradient with target network does not follow the same dynamics as SGD):**

In section 4.1, we justify why the linearisation holds for the single network case. The same reasoning does not directly hold for the target network analysis in section 4.2, and as such there is no direct theoretical guarantee that the linearisation will also hold here. We included the analogous analysis and empirical results anyway due to the relevance of the setup to practical algorithms, which we are glad the reviewer agrees are interesting negative results. We have amended the manuscript to clarify and put in better context how the analysis for the target network case is not guaranteed to approximate the true learning dynamics in the same way the single network case does.

Reviewer comment (Experiments):
The paper shows some lack of clarity by bundling together target networks with semi-gradient methods. Having a different target function is a separate design decision from choosing not to use the gradient of the target network. This suggests a possible experiment: is the linearisation better at approximating training dynamics in the semi-gradient setup without a separate target network?

Author response:
**Concern 3 (target networks and semi-gradient methods bundled together):**

We considered including an analysis for the setting suggested by the reviewer that employs semi-gradient methods without a target network. However, we decided not to because being a semi-gradient it does not satisfy the conditions for neural tangent kernel linearistation to hold (unlike the case analyzed in section 4.1) and is also not common in practice.