Simplifying Deep Temporal Difference Learning

We thank the reviewer for pointing out areas where we can improve our paper.

1. One of the strengths of PQN is indeed its sample efficiency combined with fast parallelised training. The performance of PQN is the result of several factors:

   • Removing the target networks reduces the lag in learning.
   • Adding normalisation improves stability and accelerates convergence (see Figure 6.a).
   • All the benefits of parallelisation discussed in Section 4.1.
   • To further increase sample efficiency, the use of minibatches and miniepochs allows training with experiences sampled online multiple times (see our PQN($\lambda$) algorithm).

   We realise that we did not sufficiently discuss the benefits of removing the target networks or the additional improvements in sample efficiency achievable through minibatches and miniepochs. We will add a discussion on these points in a future draft, with the aim of providing greater clarity.

2. We did conduct an ablation study to test the impact of layer normalisation on the method (Figure 6.a). We are currently working on an additional ablation to assess the effect of using multiple parallel environments with PQN, although this ablation may be challenging to perform rigorously (see the final point of our comment, “Unfair or unclear empirical evaluation,” to reviewer LgPr).

3. Many RL applications are trained in simulators, where there is no reason not to utilise parallelisation (see more on this in our comment, “PQN solves a problem different from the baseline DQN, but this was never discussed,” to reviewer LgPr). Moreover, our contribution is not solely based on parallelisation but also includes a deep theoretical analysis of the role of network normalisation in RL. Removing the replay buffer to perform parallelised sampling would make DQN very similar to PQN. However, PQN additionally employs layer normalisation to stabilise training instead of target networks, thereby improving stability and reducing the computational complexity of the algorithm.

As all of the reviewer's original concerns have been addressed in the updated draft (including new ablations) following our response below and our detailed response to Reviewer LgPr, the authors would really appreciate it if the reviewer could either acknowledge this and raise their score accordingly for the paper to be accepted if satisfied or let the authors know if they have further concerns.

Thanks so much, the authors really appreciate the feedback and the chance to improve the paper from it.

We thank the reviewer for carefully reading our work and proofs. Their feedback is really constructive and will help improve the paper further.

We have included more details on the theory in the main body of our paper. We will also provide more detail in the Appendix proofs in an updated draft once Reviewer LgPr responds to our rebuttal, we agree it can be a lot to take in on first read.

Answer to question:

We choose a unit test vector $v$ as the property $\lVert v \rVert^2=1$ makes our proofs cleaner and aligns nicely with the definition of the Matrix 2 norm, especially as we take limits of infinite network widths. As the reviewer has identified, negative definiteness on a unit circle implies negative definiteness in the whole of $\mathbb{R}^n\setminus \{0\}$: $v^\top J v<0 \implies c^2 v^\top J v<0\implies(cv)^\top J (cv)<0$ for any $c>0$, so this condition is no less general.

Thanks for the response. I of course agree that $v^T J v < 0$ for all $\| v\|=1$ is equivalent to $v^T J v$ for all $v \in \mathbb{R}^n / 0$: this was my point in my review. I can also see that restricting $\| v \| =1$ is useful in the proof. However, in the statement of the result, why not simply state that $J$ must be negative definite?

• We agree that PQN is designed to interact with an environment by taking multiple actions and observing multiple states and rewards in parallel, which may not be feasible in some scenarios. However, as the reviewer points out, it does so by building on a much simpler algorithm; removing target networks and replay through regularised networks and reverting back to the original Q-learning algorithm, with the benefit of theoretical guarentees. We are thus unsure what the review means by "probably may fail if a single environment were used." There is a wealth of empirical evidence to suggest that regularising Q-Learning [2][3][4] or actor-critic methods [5] (especially without a target network [2][5]) in single-interaction environments improves sample efficiency, computational efficiency and performance. As we state in our paper, whilst providing strong empirical evidence, this prior work offers no theoretical analysis. This is why our theoretical contribution is necessary. Repeating these experiments offers nothing in the way of contribution as we'd be repeating what has been done many times before. We will highlight this empirical evidence even further in an updated draft.
• Moreover, we wish to clarify that many RL applications are trained in simulators rather than in the real world. In these situations, there is no reason not to exploit parallelisation. Our method is the first approach that unlocks the potential of these methods through its stability, allowing for a truely online off-policy algorithm with improved sample efficieny and performance. With this in mind, DQN is a baseline for PQN because both try to apply Q-Learning to the RL problem in a sample efficient way. DQN uses a large replay buffer. PQN uses parallelised interactions. We compared them using sample efficiency as an evaluation metric. We are happy to discuss these points further in an updated daft.
• We understand the confusion in Fig. 1. We changed "Agent" for "Q-Network" and kept them separately as they can be considered two different parts of the agent (PQN keeps only the Q-Network).
• The difference between Distrubuted DQN and PQN is that the former uses multiple copies of a neural network (for instance via multi-thread or multi-machine actors) to collect experiences while continuously traning the network in a separate, parallel process (i.e. having a learner module and multiple actors modules running concurrently). PQN instead is a sequential process which involves collecting vectorised experiences, using them to train a single network. In other words, a single process running in batches. We updated the figure caption accordingly.

[2] Nauman et al, Bigger, Regularized, Optimistic: scaling for compute and sample-efficient continuous control. 2024

[3] Understanding Plasticity in Neural Networks. ICML 2023

[4] Disentangling the causes of plasticity loss in neural network. 2024

[5] CrossQ: Batch Normalization in Deep Reinforcement Learning for Greater Sample Efficiency and Simplicity. ICLR 2024

• We compare PQN mostly with PPO, which uses the same parallel environments of PQN. The comparison with non-parallelised environments is in order to compare sample efficency of PQN. Non-parallelised environments are notably much more sample efficient, and many times distributed algorithms are compared with non-parallelised only considering time as reference. On the contrary, we compare them using the number of frames, i.e. the sample efficency, where non-paralleised environments have an advantage.
• Our calculation includes parallel environments. 400M of frames is the overall number of frames collected in parallel.
• We are collecting more seeds for Atari right now and will update the figures when ready.
• DQN/Rainbow update the network every 4 frames: 200M frames/4 = 50M updates. PQN updates the network two times every 4 frames but collects 128 frames in parallel: ((200M frames/4)/128)*2 = ~70k updates
• We are open to conducting such an ablation; however, we are concerned that if done naively, it could be misleading. The issue lies in the interdependence between the number of parallel environments and other learning parameters. If we fix a budget for environment frames, collecting experiences faster generally leads to fewer network updates with larger batch sizes. This, in turn, necessitates adjusting hyperparameters such as the learning rate and the number/size of minibatches. A fair ablation would therefore require fine-tuning these hyperparameters for each number of parallel environments considered. We plan to perform such a fair ablation in at least some simpler environments, but we believe that a comprehensive analysis of this depth would be more suitable for a separate experimental study.

• Our implementation of $\lambda$-returns follows [6], where it is proved this formulation correspond to Peng's $Q(\lambda)$: "Note that the $\lambda$-return presented here unconditionally conducts backups using the maximizing action for each n-step return, regardless of which actions were actually selected by the behavioral policy μ. This is equivalent to Peng’s Q(λ)". We have added the derivation of equation 326 Q($\lambda$) in the Appendix.
• We have added Algorithm 2 in the main text. Subjectively, the authors really don't like this as we find it is messy and feel it ruins the exposition of the paper. We provided the simplest form of PQN (Algorithm 1) in the main body of the text as it is a powerful and clean expositional tool that helps readers understand our approach. Algorithm 2 is simply a generalisation of this based on $Q(\lambda)$ and we feel its details don't contribute anything to the reader's understanding. Many papers (including [6]) opt to detail algorithmic extensions they may use in the Appendix as providing these details add nothing to the paper's exposition. However, we will keep Algorithm 2 in the main body if the reviewer still disagrees with us about this.
• Our experiments follow evaluations of other convergent TD algorithms such as GTD, GTD2 and TDC (see specifically Section 11.7 and Figure 11.5 of Sutton and Barto [7]). Baird's counterexample is a provably divergent domain used to show that off-policy TD approaches diverge, that is their parameters grow without bound. Demonstrating that this is not the case, that is parameter values converge to some fixed value, is the purpose of the experiment which confirms the theoretical results. Plateauing of parameter values and thus value error is expected: like in our experiment, Figure 11.5 of [6] shows TDC learn non-zero weights and a value error that plateaus. Our results are thus no different to existing established methods.
• Finally, it is important to note from our theory that adding $\ell_2$ regularisation of magnitude $\left(\frac{\gamma}{2}\right)^2$ is just enough to ensure the TD stabilty criterion holds. We can think of this as being the edge of guaranteed stability. Without $\ell_2$ regularisation, parameters diverge in Baird's counterexample whereas with regularisation of magnitude $\left(\frac{\gamma}{2}\right)^2$, parameters converge, albeit to a plateaued value. We find it remarkable that the theory aligns so well with the empirical results. Strenghtening $\ell_2$ regularisation beyond this thus further regularises the problem, allowing us to reduce the value error to a greater extent. We have since tested and confirmed this empirically, results can be found in the updated draft.

[6] Daley, Amato, Reconciling λ-Returns with Experience Replay. 2019.

[7] Sutton and Barto, Reinforcement Learning: An Introduction. 2018

We thank the reviewer for their constructive review and appreciate that they see the strong contribution of our work. We feel we have addressed most of their concerns in the updated draft and appreciate the opportunity to improve and make our paper intelligible to as wide an audience as possible. With remaining issues, we would like to raise a few points of clarification in separate responses. As the reviewer has raised several points, we will address them here in few separate comments. Regarding the questions that the reviewer asked:

1. Yes PPO is using the same parallel environments of PQN. We've added this to the updated draft.
2. Yes we refer to Figure 12. Thank you for pointing out the typo, we've fixed it.
3. We are slightly struggling to parse the final part of this question. Do you mean what happens if an activation is placed after the final LayerNorm? If this is so, let's say the activation has Lipschitz constant $L$, then all theory remains the same except the residual term in Eq. 9 becomes $\left\lVert v_w \cdot\frac{L \gamma }{2}\right\rVert^2$. We would then need to scale the $\ell_2$ regularisation term by $L$ to compensate for this. We put the activation before the LayerNorm as this is typical for RL-specific methods like CrossQ, but can include a discussion of this in the Appendix if necessary.

1. The reviewer claims notation would be improved following Sutton and Barto. Most of our notation follows this, however when it differs it is to ensure precision and intelligibility. Firstly, Sutton and Barto do not write expectations with respect to the underlying distribution, they just use $E$ rather than $E_{x\sim P_X}$ like our notation. Whilst this may suffice for a less technical paper, the authors find it sloppy and extremely frustrating when this notation is used as the distribution is essential to determining stability of TD. Secondly, Sutton and Barto use the notation $Q(S_t,A_t)$ for $Q$-functions. We use $Q(x)$ where $x=(s,a)$ is clearly defined in the preliminaries and at several points in the paper. Switching to $Q(S_t,A_t)$ and similar notation would cause most equations to overflow onto several lines and would hinder the intelligibility of the work. Similar work proving TD convergence [1] (also authored by Sutton) uses even sparser notation, i.e. just $V^\pi$ and $V_\theta$. With this in mind, could the reviewer highlight any additional specific notation that they would wish us to change?

2. The reviewer asks for the off-policy and nonlinear parts to have their own proofs. These do have their own separate proofs that exist clearly within separate sections in the Appendix (Lemma 3 (Mitigating Off-policy Instability) and Lemma 4 (Mitigating Nonlinear Instability)). These results are summarised on separate lines in Lemma 1 of the main body. In the new draft, we have made this even clearer, labelling them off-policy bound and nonlinear bound. Aside from this, we are not sure what the reviewer is asking for here as separating the two lines into two lemmas in the main body would not add anything except take up unnecesary space in the main body. Please could the reviewer clarify what they mean?

We have added more discussion of Inqs. 3 and 4, provided more intuition of adding $\ell_2$ regularisation and a geometric interpretation of our Jacobian analysis in the updated draft, and would appreciate if the reviewer could indicate they are satisfied with this. We thank the reviewer, we feel the paper's theoretical contribution is now very simple to understand.

[1] Sutton et al., Fast Gradient-Descent Methods for Temporal-Difference Learning with Linear Function Approximation, ICML 2009

We would like to ask if the reviewer's concerns have been addressed as there is little time now to provide any more drafts before the deadline.

The authors are keen to hear back from the reviewer as to whether the latest draft (which was uploaded in colour as requested) has satisfied their concerns. There is precious little time now until the deadline and we would like to know if there are any other points that need addressing. If they are satisfied, we hope they can raise their score as promised so our work can reach a large audience.

I thank the authors for their efforts to address my concerns. The provided ablation and the author's promise to provide more runs for the Atari experiments address my concern about the empirical evaluation. Thus, I raised my score accordingly. I'm willing to increase the score further as soon as my other concerns are addressed which can be done by answering my questions and by agreeing make the necessary modifications.

I'm checking the draft at the moment. I'll leave my feedback to the parts I check as soon as I'm finish reading them. Here is my feedback on the algorithm:

There are some remaining issues in Algorithm 1:

• Using $r_t^i \sim P_R(s_t^i ,a_t^i), s_{t+1}^i \sim P_S(s_t^i ,a_t^i)$ can be misleading for the unfamiliar reader. Those distributions are not accessible for the agent. I suggest using Sutton & Barto (2018) style of writing this line "Take $a_t^i$, observe $s_{t+1}^i$ and $r_{t+1}^i$, $\forall i \in \\{0,\dots,T-1\\}$.
• The place of the $s_0\sim P_0$ is problematic. It needs to be right after "for each episode do". Also, it's missing the $i$ index for each environment.
• The notation $\\{0 : I−1\\}$ is unclear. I suggest replacing that with $\\{0,...,I−1\\}$
• I suggest the indexing starts with $1$ instead of $0$ for better readability.
• typo: mini-epochs should be epochs
• The notation $\\{t-T:t\\}$ is very ambiguous to refer to a buffer. I think the clearest way is to define a buffer $\mathcal{B}$ the is initialized to $\emptyset$ at the beginning of the algorithm. The buffer can store transitions $(s_{t}^i, a_{t}^i, r_{t+1}^i, s_{t+1}^i), \forall i \in \\{1,\dots, T \\}$ at each time step and emptied after the updating phase is complete.
• What exactly is $\pi_{\text{explore}}$? I suppose you're using $\epsilon$-greedy policy. If so, then this needs to be a clear about that and say \pi_{\text{\epsilon-greedy}}.
• The algorithm uses $x_t^i$ without a definition. I understand that it's defined in the paper, but I would be very useful for the algorithm to be self-contained such that people can understand it standalone.

Related to the algorithm:

• The recursive formulation for $\lambda$-return should have $r_{t+1}$ not $r_t$ (see line 328). The subscript in the derivation also needs to corrected as well. Additionally, I suggest using the letter $G$ instead of $R$ to avoid any confusion between return and reward. Also, lines 116 and 119, they have $r_{t}$ where it should be $r_{t+1}$.
• In line 329, $R_T$ is not defined; is it a typo? Also, the subscript in the $R_T^{\lambda}$ equation indices are not consistent with the previous line if we replaced $t$ with $T$. Lastly, it's not clear how we can end up with this equation. Can the authors provide an explanation?

• Choice of notation in Algorithm 1: We used the notation to align with [6] as well as being concise enough to fit into the page limit with all the extra information, ablations and explanations. As we have said, we can't update the draft anymore as the deadline has now passed. If it means the paper is accepted, we are more than happy to change to notation to what the reviewer has suggested.

• Exploration policy: $\pi_\textrm{explore}$ is a general exploration policy that differs from the target policy. If could be $\epsilon$-greedy, decaying $\epsilon$-greedy or it could be a more sophisticated exploration policy (based on UCB or Bayesian uncertainty estimates). The point is, for the sake of generality, alignment with theory, future research and to highlight the off-policy nature of our approach, it is very important that this is kept as general possible. We say that for our experiments we used $\epsilon$-greedy and will highlight this further when it is introduced by writing: '$\pi_\textrm{Exploration}$ is an exploration policy. We use $\epsilon$-greedy exploration in this paper, but our theory is general enough to allow for the use of more sophisticated exploration policies in future work'.

• Regards to subscripts, we are not sure that the reviewer is correct here. To see this, take $\lambda=0$, which should recover the 1-step TD update, $R_t=r_t+\gamma \max Q_\phi(s_{t+1},a')$ however under the reviewer's notation this would be $R_t=r_{t+1}+\gamma \max Q_\phi(s_{t+1},a')$.

• Yes. this is a typo, it should read $R^\lambda_t$. The term should read $R^\lambda_T=\max_{a'}Q_\phi(s_{T},a')$ which is the first starting point in the iterate process before progressing backwards. This is in line with Algorithm 1 of [6]. We will make this clearer in the final version writing:

The exploration policy $\pi_\textrm{Explore}$ is rolled out for a small trajectory of size $T$: $(s_i,a_i,r_i,s_{i+1}\dots s_{i+T})$. Starting with $R_{i+T}^\lambda=\max_{a'} Q_\phi(s_{i+T}, a')$ the targets are computed recursively back in time from $R_{i+T-1}^\lambda$ to $R_i^\lambda$ using: $R_{t}^{\lambda} =r_t + \gamma \left[ \lambda R_{t+1}^{\lambda} + (1 - \lambda) \max_{a'} Q_\phi(s_{t+1}, a') \right]$ or $R_{t}^\lambda=r_t$ if $s_{t}$ is a terminal state.

Apologies, we were really exhausted trying to get the draft ready for the original deadline. We hope there is understanding here

As this is the last day of reviewer-author discussions, we were wondering if there was anything else the reviewer is concerned about before raising their score as promised.

Thank you for the reminder. Sorry for not getting back to you sooner. I thank the authors for their reply. It is great to hear that the authors are willing to make the suggested changes if the paper gets accepted. Here is my response to some of the points:

• Re: exploration policy, I understand that you can have policies other than $\epsilon$-greedy. My point here is to have at least "e.g., $\epsilon$-greedy" as part of the algorithm to guide the reader what this exploration policy might mean.

• I suggested using $r_{t+1}$ instead of $r_t$, following Sutton and Barto (2018) since it's intuitive that the reward is given in the next step similar to the next state. It's not intuitive to have the reward and next state given in two consecutive time steps. I understand that according to your notations the definition is consistent, but I was hoping you can consider the trajectory to be $\tau_t \doteq (s_0, a_0, r_1, s_1, a_1, r_2, s_2, ..., s_{t-1}, a_{t-1}, r_t, s_t)$.

• In line 365, the authors still say "without any replay buffer". This needs to be fixed. Also, the appendix needs to be checked for same claims (e.g., we don’t use a replay buffer in line 1815)

• There is still mentioning to $Q(\lambda)$ in lines 536, 490, and 514. Also, the appendix needs to be checked for similar claims (e.g., with added normalisation and Q(λ) in line 1815)

• In line 062 there is a claim about PQN performing online updates. I think the use of the word online is not meaningful for parallel environments or with $\lambda$-return computations.

I checked the theory part once again. I think there is still room for improvement, so I encourage the authors to strive to improve accessibility for a wider range of audience. For example, a primer on Lyapunov stability needs to be added to section 2 since your analysis heavily depends on it.

In my own experience, the most confusing part was the introduction of off-policy and nonlinear components without any derivations. Now after the authors have provided the derivation in Appendix B1, the logical flow became clear and easier to follow.

There are some minor errors in the theory part:

• In line 214, you have $\delta(\phi_t) (\phi_t-\phi^\star)<0$. Are you missing a transpose to have the dot product? something like $\delta(\phi_t)^\top (\phi_t-\phi^\star)<0$?
• In assumption 2 (line 146), the TD error vector $\delta$ takes $x$ and $\phi$ as arguments. Should this be ς and $\phi$ instead? Also in line 147, should it be "Lipschitz in $\phi$,ς".
• In line 272, the authors have a limit with an index but the function doesn't have that index, so it needs to be fixed (e.g., making $Q_\phi^{\text{Layer}}$ be an explicit function of $k$.

To summarize, here are my last concerns. They are easy to fix in the final paper if the authors agreed to. Please let me know if you are willing to make these modifications if the paper gets accepted.

• Remove remaining claims about not using a replay buffer from the paper and using a more clear language instead to communicate this information. For example, "PQN uses a relatively small buffer" can replace "PQN eliminates the use of large replay buffers".

• Provide a more comprehensive discussion that the learning paradigm of PQN is heavily inspired by PPO learning paradigm in the sense you use collect data in a buffer then go over the collected data in multiple epochs and mini-batch updates then empty the buffer.

• Fixing Algorithm 1 as suggested.

• Fixing the minor errors in the theory part.

We thank the reviewer for their thoughtful and response and the time taken to review our work carefully. If space permits, depending upon Reviewer LgPr's response to our rebuttal, we will include Baird's counterexample (which we have extended to stronger $\ell_2$-regularisation) in the main body of the paper.

We thank you for your comment.

Our algorithmic focus for this paper was to develop a modern $Q$-learning approach. It is well known that $Q$-learning based algorithms are not suitable for continuous action spaces because the use of $\max_{a} Q(s',a)$.

We remark that we have provided a general and powerful theoretical analysis of TD (which applies to continuous domains), proved convergence of TD using LayerNorm + $\ell_2$ regularisation, thereby solving one of the most important open questions in RL - that is whether there exist powerful, simple nonlinear and/or off-policy TD algorithms that are provably convergent. This can stand alone as a significant theoretical contribution. In addition, we have developed a state-of-the-art $Q$-learning based algorithm that unlocks the potential of parallelised sampling. We have tested with baselines across 79 discrete-action tasks (2 Classic Control tasks, 4 MinAtar games, 57 Atari games, Craftax, 9 Smax tasks, 5 Overcooked, and Hanabi) and provided an extensive ablation study.

In contrast, the original CrossQ paper, which was an excellent piece of work and was rightly awarded a spotlight position at ICLR 2024, provided no theoretical analysis and was evaluated in only six Mujoco continuous-action tasks. For this reason, it is clear that developing a continuous actor-critic algorithm lies well beyond the scope of a conference paper, although we do hope to explore such avenues in a journal paper.

Thank you for raising this, we have immediately updated the Arxiv version with the statement `The original derivation can be found in Daley & Amato (2019, Appendix D), which we repeat and adapt here for convenience' at the start of Appendix B.4. We also request that the camera-ready revision be reopened as soon as possible be made so we can make the same change.

Our intention was never to present the derivation as original; we clearly cite Daley & Amato (2019) in the main body when referring to the algorithm and derivation. In addition, the derivation was not included in the submission - as you can see from the discussion with the reviewer below, they requested we include it for completeness. We apologise again for this oversight, and will rectify this as soon as the opportunity on open review arises. EDIT: Paper has been updated, thanks for drawing attention to this.