Simplifying Deep Temporal Difference Learning

## Abstract

Abstract Q 𝑄 Q italic_Q -learning played a foundational role in the field reinforcement learning (RL).
However, TD algorithms with off-policy data, such as Q 𝑄 Q italic_Q -learning, or nonlinear function approximation like deep neural networks require several additional tricks to stabilise training, primarily a replay buffer and target networks. Unfortunately, the delayed updating of frozen network parameters in the target network harms the sample efficiency and, similarly, the replay buffer introduces memory and implementation overheads. In this paper, we investigate whether it is possible to accelerate and simplify off-policy TD training while maintaining its stability. Our key theoretical result demonstrates for the first time that regularisation techniques such as LayerNorm can yield provably convergent TD algorithms without the need for a target network, even with off-policy data. Empirically , we find that online, parallelised sampling enabled by vectorised environments stabilises training without the need of a replay buffer. Motivated by these findings, we propose PQN, our simplified deep online Q 𝑄 Q italic_Q -Learning algorithm.
Surprisingly, this simple algorithm is competitive with more complex methods like: Rainbow in Atari, PPO-RNN in Craftax, QMix in Smax, and can be up to 50x faster than traditional DQN without sacrificing sample efficiency. In an era where PPO has become the go-to RL algorithm, PQN reestablishes off-policy Q 𝑄 Q italic_Q -learning as a viable alternative.

## 1 Introduction

In reinforcement learning (RL), the challenge of developing simple, efficient and stable algorithms remains open. Temporal difference (TD) methods have the potential to be simple and efficient, but are notoriously unstable when combined with either off-policy sampling or nonlinear function approximation (Tsitsiklis & Van Roy, ). Starting with the introduction of the seminal deep $Q$-network (DQN)(Mnih et al., ), many tricks have been developed to stabilise TD for use with deep neural network function approximators, most notably: the introduction of batched learning through a replay buffer (Mnih et al., ), target networks (Mnih et al., ), trust region based methods (Schulman et al., ), double Q-networks (Wang & Blei, ; Fujimoto et al., ), maximum entropy methods (Haarnoja et al., ; ) and ensembling (Chen et al., ). Out of this myriad of algorithmic combinations, proximal policy optimisation (PPO) (Schulman et al., ) has emerged as the de facto choice for RL practitioners, proving to be a strong and efficient baseline across popular RL domains. Unfortunately, PPO is far from stable and simple: PPO does not have provable convergence properties for nonlinear function approximation and requires extensive tuning and additional tricks to implement effectively (Huang et al., ; Engstrom et al., ).

Recent empirical studies (Lyle et al., ; ; Bhatt et al., ) provide evidence that TD can be stabilised without target networks by introducing regularisation such as BatchNorm (Ioffe & Szegedy, ) and LayerNorm (Ba et al., ; Nauman et al., ) into the $Q$-function approximator. Little is known about why these techniques work or whether they have unintended side-effects. Motivated by these findings, we ask: are regularisation techniques such as BatchNorm and LayerNorm the key to unlocking simple, efficient and stable RL algorithms? To answer this question, we provide a rigorous analysis of regularised TD. We summarise our core theoretical contributions as: I) Introduce a highly general and widely applicable analysis of TD stability;
II) we show introducing LayerNorm *and* $\ell_{2}$ regularisation into the $Q$-function approximator leads to provable convergence, stabilising nonlinear and/or off-policy TD *without the need for target networks*.

Guided by these insights, we develop a modern off-policy value-based TD method which we call a parallelised $Q$-network (PQN): for simplicity, we revisit the original $Q$-learning algorithm (Watkins, ), which updates a $Q$-function approximator online without a target network.
A recent breakthrough in RL has been running the environment and agent jointly on the GPU (Makoviychuk et al., ; Gu et al., ; Lu et al., ; Matthews et al., ; Rutherford et al., ; Lange, ). However, the replay buffer’s large memory footprint makes pure-GPU training impractical with traditional DQN. With the goal of enabling $Q$-learning in pure-GPU setting, we propose replacing the replay buffer with an online synchronous update across a large number of parallel environments. For stability, we integrate our theoretical findings in the form of a normalised deep $Q$ network. We provide a schematic of our proposed PQN algorithm in Fig. 1(d).

Figure: (a) Online $Q$-Learning
Refer to caption: x1.png

To validate our theoretical results, we evaluated PQN in Baird’s counterexample, a challenging domain that is provably divergent for off-policy methods (Baird, ). Our results show that PQN can converge where non-regularised variants fails. We provide an extensive empirical evaluation to test the performance of PQN in single-agent and multi-agent settings. Despite its simplicity, we found this algorithm competitive in a range of tasks. Notably, PQN achieves high performances in just a few hours in many games of the Arcade Learning Environment (ALE) (Bellemare et al., ), competes effectively with PPO on the open-ended Craftax task (Matthews et al., ), and stands alongside state-of-the-art Multi-Agent RL (MARL) algorithms, such as MAPPO in Overcooked (Carroll et al., ) and Hanabi (Bard et al., ) and Qmix in Smax (Rutherford et al., ). Despite not sampling from past experiences in a replay buffer, the faster convergence of PQN demonstrates that the sample efficiency loss can be minimal. This positions PQN as a strong method for efficient and stable RL in the age of deep vectorised Reinforcement Learning (DVRL).

We summarise our empirical contributions: I) we propose PQN, a simplified, parallelised, and normalised version of DQN which eliminates the use of both the replay buffer and the target network; II) we demonstrate that PQN is fast, stable, simple to implement, uses few hyperparameters, and is compatible with pure-GPU training and temporal-based networks such as RNNs, and III) our extensive empirical study demonstrates PQN achieves competitive results *in significantly less wall-clock time than existing state-of-the-art methods*.

## 2 Preliminaries

Let $\lVert\cdot\rVert$ denote the $\ell^{2}$-norm and $\mathcal{P}(\mathcal{X})$ the set of all probability distributions over $\mathcal{X}$.

### 2.1 Reinforcement Learning

In this paper, we consider the infinite horizon discounted RL setting, formalised as a Markov Decision Process (MDP) (Bellman, ; Puterman, ): $\mathcal{M}\coloneqq\langle\mathcal{S},\mathcal{A},P_{S},P_{0},P_{R},\gamma\rangle$ with bounded state space $\mathcal{S}$, bounded action space $\mathcal{A}$, transition distribution $P_{S}:\mathcal{S}\times\mathcal{A}\to\mathcal{P}(\mathcal{S})$, initial state distribution $P_{0}\in\mathcal{P}(\mathcal{S})$, bounded stochastic reward distribution $P_{R}:\mathcal{S}\times\mathcal{A}\to\mathcal{P}([-r_{max},r_{max}])$ where $r_{max}\in\mathbb{R}<\infty$ and scalar discount factor $\gamma\in[0,1)$. An agent in state $s_{t}\in\mathcal{S}$ taking action $a_{t}\in\mathcal{A}$ observes a reward $r_{t}\sim P_{R}(s_{t},a_{t})$. The agent’s behaviour is determined by a policy that maps a state to a distribution over actions: $\pi:\mathcal{S}\to\mathcal{P}(\mathcal{A})$ and the agent transitions to a new state $s_{t+1}\sim P_{S}(s_{t},a_{t})$. As the agent interacts with the environment through a policy $\pi$, it follows a trajectory $\tau_{t}\coloneqq(s_{0},a_{0},r_{0},s_{1},a_{1},r_{1},\dots s_{t-1},a_{t-1},r_
{t-1},s_{t})$ with distribution $P^{\pi}_{t}$. For simplicity, we denote state-action pair $x_{t}\coloneqq(s_{t},a_{t})\in\mathcal{X}$ where $\mathcal{X}\coloneqq\mathcal{S}\times\mathcal{A}$. The state-action pair transitions under policy $\pi$ according to the distribution $P_{X}^{\pi}(x):\mathcal{X}\rightarrow\mathcal{P}(\mathcal{X})$.

The agent’s goal is to learn an optimal policy of behaviour $\pi^{\star}\in\Pi^{\star}$ by optimising the expected discounted sum of rewards over all possible trajectories $J^{\pi}$, where: $\Pi^{\star}\coloneqq\operatorname*{arg\,max}_{\pi}J^{\pi}$ is an optimal policy for the objective $J^{\pi}\coloneqq\mathbb{E}_{\tau_{\infty}\sim P_{\infty}}\left[\sum_{t=0}^{
\infty}\gamma^{t}r_{t}\right]$. The expected discounted reward for an agent in state $s_{t}$ for taking action $a_{t}$ is characterised by a $Q$-function, which is defined recursively through the Bellman equation: $Q^{\pi}(x_{t})=\mathcal{B}^{\pi}[Q^{\pi}](x_{t})$, where the Bellman operator $\mathcal{B}^{\pi}$ projects functions forwards by one step through the dynamics of the MDP:

$$ $\displaystyle\mathcal{B}^{\pi}[Q^{\pi}](x_{t})\coloneqq\mathbb{E}_{x_{t+1}\sim P _{X}^{\pi}(x_{t}),r_{t}\sim P(x_{t})}\left[r_{t}+\gamma Q^{\pi}(x_{t+1})\right].$ (1) $$

Of special interest is the $Q$-function for an optimal policy $\pi^{\star}$, which we denote as $Q^{\star}(x_{t})\coloneqq Q^{\pi_{\star}}(x_{t})$. The optimal $Q$-function satisfies the optimal Bellman equation $Q^{\star}(x_{t})=\mathcal{B}^{\star}[Q^{\star}](x_{t})$, where $\mathcal{B}^{\star}$ is the optimal Bellman operator:

$$ $\displaystyle\mathcal{B}^{\star}[Q^{\star}](x_{t})\coloneqq\mathbb{E}_{s_{t+1} \sim P_{S}(x_{t}),r_{t}\sim P(x_{t})}\left[r_{t}+\gamma\max_{a^{\prime}}Q^{ \star}(s_{t+1},a^{\prime})\right].$ (2) $$

### 2.2 Temporal Difference Methods

Many RL algorithms employ TD learning for policy evaluation, which combines bootstrapping, state samples and sampled rewards to estimate the expectation in the Bellman operator (Sutton, ). We introduce a $Q$-function approximation $Q_{\phi}:\mathcal{X}\to\mathbb{R}$ parametrised by $\phi\in\Phi$ to represent the space of $Q$-functions. We assume that $Q_{\phi}$ is initialised from a distribution $\phi_{0}\sim P_{\Phi}$. In their simplest form, TD methods estimate the application of a Bellman operator by updating the $Q$-function approximator parameters according to:

$$ $\displaystyle\phi_{i+1}=\phi_{i}+\alpha_{i}\left(r+\gamma Q_{\phi_{i}}(x^{ \prime})-Q_{\phi_{i}}(x)\right)\nabla_{\phi}Q_{\phi_{i}}(x),$ (3) $$

where $x\sim d^{\mu},r\sim P_{R}(x),x^{\prime}\sim P_{X}^{\pi}(x)$ and $\alpha_{i}$ is a sequence of stepsizes satisfying the standard Robbins-Munro conditions (Robbins & Monro, ):

###### Assumption 1 (RM Conditions) .

We assume $\alpha_{i}>0$ with $\sum_{i=0}^{\infty}\alpha_{i}=\infty$ and $\sum_{i=0}^{\infty}\alpha_{i}^{2}<\infty$.

Here $d^{\mu}\in\mathcal{P}(\mathcal{X})$ is a sampling distribution, and $\mu$ is a sampling policy that may be different from the target policy $\pi$. Methods for which the sampling policy differs from the target policy are known as *off-policy* methods. In this paper, we will study the $Q$-learning (Watkins, ; Dayan, ) TD update: $\phi_{i+1}=\phi_{i}+\alpha_{i}(r+\gamma\sup_{a^{\prime}}Q_{\phi_{i}}(s^{\prime
},a^{\prime})-Q_{\phi_{i}}(x))\nabla_{\phi}Q_{\phi_{i}}(x)$, which aims to learn an optimal $Q$-function by estimating the optimal Bellman operator. As data in $Q$-learning is gathered from an exploratory policy $\mu$ that is not optimal, $Q$-learning is an inherently off-policy algorithm. For simplicity of notation we define the tuple $\varsigma\coloneqq(x,r,x^{\prime})$ with distribution $P_{\varsigma}$ and the TD-error vector as:

$$ $\displaystyle\delta(\phi,\varsigma)\coloneqq\left(r+\gamma Q_{\phi}(x^{\prime} )-Q_{\phi}(x)\right)\nabla_{\phi}Q_{\phi}(x),$ (4) $$

allowing us to write the TD parameter update as: $\phi_{i+1}=\phi_{i}+\alpha_{i}\delta(\phi_{i},\varsigma)$. Typically, $d^{\mu}$ is the stationary state-action distribution of an ergodic Markov chain but may be another offline distribution such as a distribution induced by a replay buffer. We introduce the following mild regularity assumptions our analysis.

###### Assumption 2 (Regularity Assumptions) .

Assume that $\Phi\subset\mathbb{R}^{d}$ is compact and convex and $\delta(x,\phi)$ is Lipschitz in $\phi,s$. When updating TD, $x\sim d^{\mu}$ is either sampled i.i.d. from a distribution with support over $\mathcal{X}$ or is sampled from a geometrically ergodic Markov chain with stationary distribution $d^{\mu}$ defined over a $\sigma$-algebra that is countably generated from $\mathcal{X}$.

The condition of $\Phi\subset\mathbb{R}^{d}$ being compact is ubiquitous in TD theory and stochastic approximation (Papavassiliou & Russell, ; Nemirovski et al., ; Maei et al., ; Kushner, ; Lacoste-Julien et al., ; Bhandari et al., ; Wang et al., ; Yang et al., ; Zhang et al., ) and can be achieved by projecting any $\phi\notin\Phi$ back into $\Phi$ using the projection $P_{\Phi}(\phi^{\prime})\coloneqq\operatorname*{arg\,min}_{\phi}\in\Phi(\lVert
\phi-\phi^{\prime}\rVert)$. Projection is a mathematical formality and should not be required in practice as $\Phi$ can be made large enough to contain all updates when TD is stable and a suitable stepsize regime is chosen. Finally, geometric ergodicity extends traditional notions of aperiodicity and irreducibility in discrete MDPs to the more general continuous state-action space formulations (see Roberts & Rosenthal () for details). It is one of the weakest ergodicity assumptions.

We denote the expected TD-error vector as:

$$ $\displaystyle\delta(\phi)\coloneqq\mathbb{E}_{\varsigma\sim P_{\varsigma}}[ \delta(\phi,\varsigma)],$ (5) $$

and define the set of TD fixed points as: $\phi^{\star}\in\left\{\phi|\delta(\phi)=0\right\}$. If a TD algorithm converges, it must converge to a TD fixed point as the expected parameter update is zero for all $\phi^{\star}$.

###### Assumption 1 (RM Conditions) .

###### Assumption 2 (Regularity Assumptions) .

### 2.3 Vectorised Environments

Parallelising the interactions between an RL agent and a learning environment is a standard method for speeding up training. In classical frameworks like Gymnasium (Towers et al., ), this is achieved by processing multiple environments via multi-threading. In more recent GPU-based frameworks like IsaacGym (Makoviychuk et al., ), ManiSkill2 (Gu et al., ), Jumanji (Bonnet et al., ), Craftax (Matthews et al., ), JaxMARL (Rutherford et al., ) the environments’ operations are vectorised, meaning that they are performed together using batched tensors. This allows an agent to easily interact with thousands of environments, and it enables the compilation of end-to-end GPU learning pipelines, which can accelerate the training of on-policy agents like PPO and A2C by orders of magnitude (Makoviychuk et al., ; Weng et al., ; Gu et al., ; Lu et al., ). Unfortunately, end-to-end single-GPU training is not compatible with traditional off-policy methods like DQN for two reasons: firstly, maintaining a replay buffer in GPU is not feasible in complex environments, as it would occupy most of the GPU memory; and secondly, the convergence of off-policy methods demands a very high number of updates in relation to the sampled experiences (DQN traditionally performs one gradient step per environment step). Commonly, parallelisation of Q-Learning (like in Ape-X (Horgan et al., ), R2D2 (Kapturowski et al., ) and a recent method presented in Li et al. ()) is achieved by continuously training the Q-network in a separate process in order to keep up with the fast sampling (see Fig. 1(c)), a setup that is not feasible in a single pure-GPU setting. For this reason, all referenced frameworks primarily provide PPO or A2C baselines, i.e. vectorised RL lacks a off-policy $Q$-learning baseline.

### 2.4 Related Work

Our paper makes several significant contributions across a range of interconnected threads in RL research, which prohibits a full discussion of all related methods in the main body. We provide an extensive discussion of all related work in Appendix A.

## 3 Analysis of Regularised TD

*Proofs for all theorems and corollaries can be found in Appendix B*

Building on (Bhandari et al., ; Fellows et al., ), we now develop a powerful and general Jacobian analysis tool to characterise stability of TD approaches used in practice (Section 3.1). We then apply this
analysis to regularised TD, confirming our theoretical hypothesis that careful application of LayerNorm and $\ell_{2}$ regularisation can stabilise TD (Section 3.2). Finally, we compare LayerNorm to BatchNorm regularisation techniques in Section 3.3, explaining our preference for LayerNorm. Recalling that $x\coloneqq\{s,a\}$, we remark that our results can be derived for value functions by setting $x=s$ in our analysis.

### 3.1 Stability of TD

As TD updates aren’t a gradient of any objective, they fall under the more general class of algorithms known as stochastic approximation (Robbins & Monro, ; Borkar, ). Stability is not guaranteed in the general case and convergence of TD methods has been studied extensively (Watkins & Dayan, ; Tsitsiklis & Van Roy, ; Dalal et al., ; Bhandari et al., ; Srikant & Ying, ). We now extend the methods of Fellows et al. () to include more general nonlinear function approximators (including neural networks with Lipschitz activations such as ReLU), meaning our analysis applies exactly to TD methods used in practice. Key to determining the stability of the TD is establishing the negative definiteness of the TD Jacobian.

If this criterion holds, then as stepsizes approach zero in the limit $\lim_{i\rightarrow\infty}\alpha_{i}=0$, there exists some $t$ such that for every $i>t$ each update is a contraction mapping: $\mathbb{E}\left[\left\lVert\phi_{i+1}-\phi^{\star}\right\rVert\right]<\mathbb{
E}\left[\left\lVert\phi_{i}-\phi^{\star}\right\rVert\right]$. This key condition allows us to prove convergence of TD:

###### Theorem 1 (TD Stability) .

Let Assumptions 1 and 2 hold. If the TD criterion holds then the TD updates in Eq. 3 converge with: $\lim_{i\rightarrow\infty}\mathbb{E}\left[\left\lVert\phi_{i}-\phi^{\star}
\right\rVert^{2}\right]=0$.

Analysing the TD Jacobian reveals two inherent causes of instability:

Off-policy Instability: TD can become unstable if:

$$ $\displaystyle\gamma v^{\top}\mathbb{E}_{\varsigma\sim P_{\varsigma}}\left[ \nabla_{\phi}Q_{\phi}(x^{\prime})\nabla_{\phi}Q_{\phi}(x)^{\top}\right]v<v^{ \top}\mathbb{E}_{x\sim d^{\mu}}\left[\nabla_{\phi}Q_{\phi}(x)\nabla_{\phi}Q_{ \phi}(x)^{\top}\right]v,$ (6) $$

does not hold for any test vector $v$. Unfortunately, this condition can only be shown to hold in the on-policy sampling regime when $d^{\mu}=d^{\pi}$ is the stationary distribution of the MDP. This is because the Cauchy-Schwartz inequality needed to prove $\gamma\lVert\nabla_{\phi}Q_{\phi}(x^{\prime})\nabla_{\phi}Q_{\phi}(x)^{\top}\|
\leq\lVert\nabla_{\phi}Q_{\phi}(x)\nabla_{\phi}Q_{\phi}(x)^{\top}\|$ does not hold for arbitrary off-policy sampling distributions $d^{\mu}$ (Fellows et al., ). Hence *off-policy sampling* is a key source of instability in TD, especially in algorithms such as $Q$-learning.

Nonlinear Instability: TD can become unstable if:

$$ $\displaystyle\mathbb{E}_{\varsigma\sim P_{\varsigma}}\left[(r+\gamma Q_{\phi}( x^{\prime})-Q_{\phi}(x))v^{\top}\nabla_{\phi}^{2}Q_{\phi}(x)v\right]<0,$ (7) $$

does not hold for any unit test vector $v$. This condition does not apply in the linear case as second order derivatives are zero. In the nonlinear case, the left hand side can be arbitrarily positive depending upon the specific MDP and choice of function approximator. Hence *nonlinearity* is a key source of instability in TD. Together both off-policy and nonlinear instability formalise the *deadly triad* (Sutton & Barto, ; van Hasselt et al., ) and TD can be unstable if *either* Conditions 6 or 7 are not satisfied. We now investigate how LayerNorm with $\ell_{2}$ regularisation can tackle these sources of instability.

###### Theorem 1 (TD Stability) .

### 3.2 Stabilising TD With LayerNorm + ℓ 2 subscript ℓ 2 \ell_{2} roman_ℓ start_POSTSUBSCRIPT 2 end_POSTSUBSCRIPT Regularisation

To understand how LayerNorm with $\ell_{2}$ regularisation stabilises TD, we study the following $Q$-function approximator:

$$ $\displaystyle\vspace{-1cm}Q_{\phi}^{\textrm{Layer}}(x)=w^{\top}\textrm{ LayerNorm}^{k}\left[\sigma\circ Mx\right].$ (8) $$

Here $\phi=[w^{\top},\textrm{Vec}(M)^{\top}]$ is the parameter vector where $M\in\mathbb{R}^{k\times d}$ is a $k\times d$ matrix, $\sigma$ is an element-wise Lipschitz activation and $w\in\mathbb{R}^{k}$ is a vector of final layer weights. The LayerNorm (Ba et al., ) activation is defined element-wise as:

$$ $\displaystyle\textrm{LayerNorm}^{k}_{i}\left[f(x)\right]$ $\displaystyle\coloneqq\frac{1}{\sqrt{k}}\frac{f_{i}(x)-\frac{1}{k}\sum_{j=0}^{ k-1}f_{j}(x)}{\sqrt{\sum_{i=0}^{k-1}(f_{i}(x)-\frac{1}{k}\sum_{j=0}^{k-1}f_{j} (x))^{2}+\epsilon}},$ (9) $$

where $\epsilon>0$ is a small constant introduced for numerical stability. Deeper networks with more LayerNorm layers may be used in practice, however our analysis reveals that only the final layer weights affect the stability of TD with wide LayerNorm neural networks. We observe that adding LayerNorm does not affect the representational capacity of the network as it merely rescales the input according to a standard Gaussian. The output is then rescaled due to the final linear layer. As $K$ increases, the empirical mean and standard deviations in Eq. 9 approach their true expectations, thereby increasing the degree of normalisation provided. We now investigate theoretically what effect this has on the off-policy and nonlinear stability conditions introduced in Section 3.1:

###### Lemma 1 .

Let 2 apply. Using the LayerNorm $Q$-function defined in Eq. 8:

$$ $\displaystyle\gamma v^{\top}\nabla_{\phi}Q_{\phi}(x^{\prime})\nabla_{\phi}Q_{ \phi}(x)^{\top}v-(v^{\top}\nabla_{\phi}Q_{\phi}(x))^{2}$ $\displaystyle\leq\left(\lVert v\rVert\nicefrac{{\gamma}}{{2}}\right)^{2}+ \mathcal{O}\left(\nicefrac{{1}}{{\sqrt{k}}}\right),$ (10) $\displaystyle\left\lvert\left(r+\gamma Q_{\phi}(x^{\prime})-Q_{\phi}(x)\right) v^{\top}\nabla_{\phi}^{2}Q_{\phi}(x)v\right\rvert$ $\displaystyle=\mathcal{O}\left(\nicefrac{{1}}{{\sqrt{k}}}\right),$ (11) $$

almost surely for any test vector and any state-action transition pair $x,x^{\prime}\in\mathcal{X}$.

Eq. 11 of Lemma 1 reveals that all nonlinear instability can be mitigated as $k\rightarrow\infty$ hence for a suitably wide critic the only form of instability that we need to account for is the off-policy instability (the $(\nicefrac{{\gamma L_{\sigma}}}{{2}})^{2}$ term in Inq. 10) *and only in the final layer weights*. Crucially, the bound in Eq. 10 *does not depend on the sampling distribution* $d^{\mu}$. Mathematically, this is due to the normalising property: $\lVert\textrm{LayerNorm}[Mx]\rVert\leq 1$, which is independent of the input data. In contrast, a wide neural network without LayerNorm would tend towards a linearised form $Q_{\phi_{0}}(x)+\nabla_{\phi}Q_{\phi_{0}}(x)^{\top}(\phi-\phi_{0})$ (Lee et al., ) which is data dependent and still suffers from off-policy instability like all linear function approximators in the off-policy regime. To incorporate $\ell_{2}$ regularisation, we propose regularising the final weight layer $w$ using the following TD update vector:

$$ $\displaystyle\delta^{\textrm{reg}}(\phi,\varsigma)\coloneqq\delta^{\textrm{ Layer}}(\phi,\varsigma)-\left(\frac{\gamma}{2}\right)^{2}[w^{\top}\ 0\ \cdots 0 ]^{\top},$ (12) $$

where $\delta^{\textrm{Layer}}(\phi,\varsigma)$ is the TD update vector from Eq. 4 using the LayerNorm critic from Eq. 8 respectively. We now formally verify that given sufficiently large $k$, regularised TD updates can satisfy the TD stability criterion, thereby stabilising TD:

###### Theorem 2 .

Let 2 apply. Using the LayerNorm regularised TD update in Eq. 12, there exists some finite $k^{\prime}$ such that the TD stability criterion holds for all $k>k^{\prime}$.

In Section 5.1 we test our theoretical claim in Theorem 2 empirically, demonstrating that LayerNorm + $\ell_{2}$ regularisation can stabilise Baird’s counterexample, an MDP intentionally designed to cause instability for TD (Baird, ).

Finally, we remark that whilst adding an $\ell_{2}$ regularisation term $-\eta\phi$ to all parameters can stabilise TD in principle, large $\eta$ recovers a quadratic optimisation problem with minimum at $\phi=0$, pulling the TD fixed points towards 0 0. Hence, we suggest $\ell_{2}$-regularisation should be used sparingly; only when LayerNorm alone cannot stabilise the environment and only over the final layer weights. Aside from Baird’s counterexample, we find LayerNorm without $\ell_{2}$ regularisation can stabilise all environments in our extensive empirical evaluation in Section 5.

###### Lemma 1 .

###### Theorem 2 .

### 3.3 LayerNorm and BatchNorm TD

We have seen from Theorem 2 that LayerNorm + $\ell_{2}$ regularised TD can stabilise TD by mitigating the effects of nonlinearity and off-policy sampling. Empirical evidence suggests that BatchNorm (Ioffe & Szegedy, ) regularisation, which is essential for stabilising algorithms such as CrossQ (Bhatt et al., ), may also possess similar properties to LayerNorm. It is natural to ask: ‘what are the potential benefits of LayerNorm over BatchNorm methods?’

It is known empirically that naïvely applying BatchNorm as presented by Ioffe & Szegedy () does not stabilise TD as CrossQ does not succeed without applying several modifying ‘tricks’ such as double Q-learning, batch renormalisation using running statistics and calculating the batch statistics from a mixture of datasets (Bhatt et al., ). This creates a theoretical gap between analysis of the BatchNorm function and what is applied in practice. In contrast, LayerNorm can be applied out of the box to stabilise RL, hence the implementation does not deviate from the theoretical analysis, maintaining the strong theoretical guarantees in Theorem 2. A practical algorithm can be obtained using LayerNorm without burdening practioners with additional tricks and their associated hyperparameter tuning. Additionally, compared to BatchNorm, LayerNorm is simpler to implement and does not require memory or estimation of the running batch averages.

Our empirical analysis in Section 5 shows that BatchNorm can degrade performance in some cases, while in others it can improve results if applied early in the network. Therefore, we don’t dismiss BatchNorm outright, but a thorough theoretical analysis is needed to fully understand its practical effects. Nonetheless, we recommend starting with LayerNorm and $\ell_{2}$ regularisation as a strong, simple baseline for stabilising TD algorithms before experimenting with alternatives like BatchNorm.

## 4 Parallelised Q 𝑄 Q italic_Q -Learning

Guided by the analysis in Section 3, we develop a simplified version of deep $Q$-learning

Figure: Algorithm 1 PQN ($s,a,s^{\prime},r,\textrm{target}$ are $B$-dim vectors)

to exploit the power of parallelised sampling without a replay buffer or target networks. The Q-Network is regularised with network normalisation (preferably LayerNorm) and $\ell_{2}$ regularisaiton as required (see Eq. 12). We present our algorithm in Algorithm 1. PQN is equivalent to a vectorised variant of Watkins ()’s original $Q$-learning algorithm where $B$ separate interactions occur in parallel with the environment and a normalised network is used as function approximation.

Table 1 summarises the advantages of PQN in comparison to popular methods. Compared to traditional DQN and distributed DQN, PQN enjoys ease of implementation, fast execution, very low memory requirements, and high compatibility with GPU-based training and RNNs. The only algorithm that shares these attributes is PPO.
However, although PPO is in principle a simple algorithm, its success is determined by numerous interacting implementation details (Huang et al., ; Engstrom et al., ), making the actual implementation challenging. Moreover, PQN uses few main hyperparameters, namely the number of parallel environments, the learning rate and epsilon with its decay.

As PQN is an online algorithm, it is straightforward to implement a version that uses $n$-step returns. In our implementation, we experiment with Peng’s $Q(\lambda)$ (Peng & Williams, ) by rolling the policy for a small trajectory of size $T$ and then computing the targets recursively with the observed rewards and the Q-values computed by the agent at the next state: $R_{t}^{\lambda}=r_{t}+\gamma\left[\lambda R_{t+1}^{\lambda}+(1-\lambda)\max_{a
^{\prime}}Q(s^{\prime},a^{\prime})\right]$, where $R_{T}^{\lambda}=r_{t}+\gamma\max_{a^{\prime}}Q(s^{\prime},a^{\prime})$. Similarly, an implementation of PQN using RNNs only requires sampling trajectories for multiple time-steps and then back-propagating the gradient through time in the learning phase. This is much easier than SOTA $Q$-Learning algorithms based on replay buffer that use $Q(\lambda)$ or RNNs (such as Retrace (Munos et al., ) and R2D2), which adopt computationally intensive techniques to handle the computation of $Q(\lambda)$ and the hidden states of trajectories collected with old policies during replay (see Appendix C for details). Finally, a basic multi-agent version of PQN for coordination problems can be obtained by adopting Value Network Decomposition Networks (VDN) (Sunehag et al., ), i.e. optimising the joined action-value function as a sum of the single agents action-values.

**Table 1: Advantages and Disadvantages of DQN, Distributed DQN, PPO and PQN.**
|  | DQN | Distr. DQN | PPO | PQN |
| --- | --- | --- | --- | --- |
| Implementation | Easy | Difficult | Medium | Very Easy |
| Memory Requirement | High | Very High | Low | Low |
| Training Speed | Slow | Fast | Fast | Fast |
| Sample Efficient | Yes | No | Yes | Yes |
| Compatibility with RNNs | Medium | Medium | High | High |
| Compatibility w. end-to-end GPU Training | Low | Low | High | High |
| Amount of Hyper-Parameters | Medium | High | Medium | Low |
| Convergence | No | No | No | Yes |

### 4.1 Benefits of Online Q 𝑄 Q italic_Q -learning with Vectorised Environments

Figure: Figure 2: Sketch of Sampling Regimes in DQN and PQN
Refer to caption: x5.png

Compared to a replay buffer, vectorisation of the environment enables fast collection of many parallel transitions from independent trajectories.
Denoting the stationary distribution at time $t$ of the MDP under policy $\pi_{t}$ as $d^{\pi_{t}}$, uniformly sampling from a replay buffer estimates sampling from the average of all distributions across all timesteps: $\frac{1}{t^{\prime}+1}\sum_{t=0}^{t^{\prime}}d^{\pi_{t}}$. In contrast, vectorised sampling in PQN estimates sampling from the stationary distribution $d^{\pi_{t^{\prime}}}$ at timestep $t$’. We sketch the difference in these sampling regimes in Fig. 2. Coloured lines represent different state-actions trajectories across the vectorised environment as a function of timestep $t$. Crosses represent samples drawn for each algorithm at timestep $t^{\prime}$.

PQN’s sampling further aids algorithmic stability by better approximating this regime in two ways: firstly, the parallelised nature can help exploration since the (potential) natural stochasticity in the dynamics means even a greedy policy will explore several different states in parallel. Secondly, by taking multiple actions in multiple states, PQN’s sampling distribution is a good approximation of the true stationary distribution under the current policy: as time progresses, ergodic theory states that this sampling distribution converges to $d^{\pi_{t^{\prime}}}$. In contrast, sampling from a replay buffer involves sampling from an average of *older* stationary distributions under shifting policies from a single agent, which will be more offline and take longer to converge, as illustrated in Fig. 2. We emphasise that *PQN is still an off-policy approach* since it uses two different policies to optimise the Bellman equations: the $\epsilon$-greedy policy for the current timestep and the current policy for the next. Notice that at beginning of training PQN uses an $\epsilon=1$, meaning that it approximates a value function from a completely random policy. This requires normalisation to mitigate off-policy instability identified in Section 3.

## 5 Experiments

In contrast to prior work in $Q$-learning, which has focused heavily on evaluation in the Atari Learning Environment (ALE) (Bellemare et al., ), probably overfitting to this environment, we evaluate PQN on a range of single- and multi-agent environments, with PPO as the primary baseline. We summarise the memory and sample efficiency of PQN in Table 2. Due to our extensive evaluation, additional results are presented in Appendix E. All experimental results are shown as mean of 10 seeds, except in ALE where we followed a common practice of reporting 3 seeds.

### 5.1 Confirming Theoretical Results

Fig. 6 (a) of Appendix E shows that together LayerNorm + $\ell_{2}$ can stabilise TD in Baird’s counterexample (Baird, ), a challenging environment that is intentionally designed to be provably divergent, even for linear function approximators. Our results show that stabilisation is mostly attributed to the introduction of LayerNorm. Moreover the degree of $\ell_{2}$-regularisation needed is small - just enough to mitigate off-policy stability due to final layer weights according to Theorem 2 - and it makes relatively little difference when used in isolation.

### 5.2 Atari

Figure: (a) Atari-10 Score
Refer to caption: x6.png

To save computational resources, we evaluate PQN against PPO in the Atari-10 suite of games from the ALE, which estimates the median across the full suite using a smaller sample of games. PQN outperforms PPO in terms of sample efficiency, final score, and training time (1 hour compared to 2.5 hours for PPO), and also surpasses sample-efficient methods like Double-DQN and Prioritised DDQN in the same number of frames, despite these methods being trained for several days and using over 50 times more gradient updates (50M compared to 700k for PQN). To further test our method, we train PQN on the full suite of 57 Atari games. Fig. 3(d) shows that the time needed to train PQN on the full Atari suite is equivalent to the time required to train traditional DQN methods on a single game(^1^11DQN training time was optimistically estimated using the JAX-based CleanRL DQN implementation.). With an additional budget of 100M frames (30 minutes of training), PQN achieves the median score of Rainbow (Hessel et al., ), which is still a SOTA method in ALE for sample efficiency but requires around 3 days of training per game, meaning that PQN can be considered 50x faster. While Rainbow is slightly more sample efficient, it’s important to note that Rainbow is a much more complex system, designed specifically for Atari. Moreover, parallelisation of $Q$-Learning has traditionally sacrificed far more sample efficiency than PQN. For instance, Ape-X struggles to solve even the simplest Atari game, Pong, within 200M frames (Horgan et al., ). In this regard, PQN represents a significant advancement in $Q$-Learning research, offering a balanced compromise between speed, simplicity, and sample efficiency.

In Appendix E, we provide detailed data from these experiments, a comparison with Dopamine-Rainbow using the IQM score, and a comparative histogram of the performances of algorithms in all the games. In this histogram, we show that PQN reaches human-level performance in 40 of the 57 games of the ALE, underperforming mainly in the hard-exploration games, suggesting that the $\epsilon$-greedy exploration used by PQN is too simple to solve ALE, and indicating a clear research direction to improve the method. Moreover, the absence of a replay buffer precludes the use of a prioritised buffer, which can be viewed as a limitation of PQN but also as a challenge for future DVRL research.

### 5.3 Craftax

Craftax (Matthews et al., ) is an open-ended RL environment based on Crafter (Hafner, ) and Nethack (Küttler et al., ). It is a challenging environment requiring solving multiple tasks to be completed. By design, Craftax is fast to run in a pure-GPU setting, but existing benchmarks are based solely on PPO. The observation size of the symbolic environment is around 8000 floats, making a pure-GPU DQN implementation with a buffer prohibitive, as it would take around 30GBs of GPU-ram. PQN can provide an off-policy $Q$-learning baseline without using GPU memory for a replay buffer. Following the Craftax paper, we evaluate for 1B steps and compared PQN to PPO using both an MLP and an RNN. The RNN results are shown in Fig. 4(a). PQN is more sample efficient and with a RNN obtains a higher score of 16% against the 15.3% of PPO-RNN. The two methods also take a similar amount of time to train. We consider it a success to offer researchers a simple $Q$-learning alternative to PPO that can run on GPU on this challenging environment.

### 5.4 Multi-Agent Tasks

Figure: (a) Craftax
Refer to caption: x10.png

When dealing with multi-agent problems, the replay buffer needs to store observations for all agents, increasing the memory requirements up to hundreds of gigabytes. Additionally, RNNs are highly effective in handling the individual agents’ partial observability of the environments, and credit assignment, a key challenge in MARL, is typically addressed with value-based methods Sunehag et al. (); Rashid et al. (). Therefore, a memory-efficient, RNN-compatible, and value-based method is highly desirable. We evaluate PQN combined with VDN in Hanabi (Bard et al., ), SMAC-SMACV2 (Ellis et al., ; Samvelyan et al., ) (in its JAX-vectorised version, Smax)  (Rutherford et al., ), and Overcooked (Carroll et al., ). Smax is a faster version of SMAC, running entirely on a single GPU. Notably, when at least 20 agents are active in the environment, a replay buffer can consume all available memory on a typical 10GB GPU. PQN-VDN runs successfully on Smax without a replay buffer, outperforming MAPPO and QMix. Remarkably, PQN learns a coordination policy even in the most difficult scenarios in about 10 minutes, compared to QMix’s 1 hour (see Fig. 15). Similarly, PQN outperforms the replay-buffer-based version of VDN and PPO in Overcooked, and is significantly more sample-efficient than MAPPO in Hanabi, where it achieves an average score of 24 points.

### 5.5 Ablations

To examine the effectiveness of PQN’s algorithmic components, we perform the following ablations.

#### Regularisation:

In Fig. 5(a), we examine the impact of regularisation on performance in the Atari-10 suite. Results show that LayerNorm significantly improves performance, supporting the theoretical findings in Section 3, while BatchNorm can degrade performance. Additionally, applying the additional tricks from CrossQ further worsens PQN’s performance.

#### Input Normalisation:

In preliminary experiments, we observed that BatchNorm significantly improves PQN performances in Craftax. Figure 5(b) compares the performance of PQN RNN with BatchNorm, LayerNorm, and no normalisation in the two cases where BatchNorm is applied to the input before the first hidden layer or not. Without input normalisation, BatchNorm provides a substantial boost. However, PQN performs best when only the input to the first layer is normalised, and applying LayerNorm to the rest of the network offers a similar improvement. This suggests BatchNorm can be effective as input normalisation, particularly in scenarios like Craftax with large, sparse observation vectors.

#### Varying λ 𝜆 \lambda italic_λ :

In Fig. 5(c) we compare different values of $\lambda$ in Atari-10. We find that a value of $\lambda=0.65$ performs the best by a significant margin. It significantly outperforms $\lambda=0$ (which is equal to perform 1-step update with the traditional Bellman operator) confirming that the use of $Q(\lambda)$ represents an important design choice over one-step TD.

#### Replay Buffer:

In Fig. 5(d), we compare PQN from Algorithm 1 with a variant that maintains a standard sized replay buffer of 1M of experiences in GPU using Flashbax Toledo et al. ().
This version converges to the same final performance but takes $\sim$6x longer to train, which is likely due to the constant need to perform random access of a buffer of around 30GBs. This reinforces our core message that a large memory buffer should be avoided in pure GPU training.

Figure: (a) Atari-10 Normalisation
Refer to caption: x14.png

## 6 Conclusion

**Table 2: Summary of Memory Saved and Speedup of PQN Compared to Baselines. The Atari speedup is relative to the traditional DQN pipeline, which runs a single environment on the CPU while training the network on GPU. Smax and Craftax speedups are relative to baselines that also run entirely on GPU but use a replay buffer. The Hanabi speed-up is relative to an R2D2 multi-threaded implementation.**
|  | Memory Saved | Speedup |
| --- | --- | --- |
| Atari | 26gb | 50x |
| Smax | 10gb (up to hundreds) | 6x |
| Hanabi | 250gb | 4x |
| Craftax | 31gb | 6x |

We have presented the first rigorous analysis explaining the stabilising properties of LayerNorm and $\ell_{2}$ regularisation in TD methods. These results allowed us to develop PQN, a simple, stable and efficient regularised $Q$-learning algorithm without the need for target networks or a replay buffer. PQN exploits vectorised computation to achieve excellent performance across an extensive empirical evaluation with a significant boost in computational efficiency, without sacrificing sample efficiency. PQN offers a simple pipeline that is easy to implement and out-of-the-box compatible with key elements in RL, such as $Q(\lambda)$ and RNNs, which are otherwise difficult to use in current $Q$-Learning implementations. Additionally, it provides a valuable baseline for multi-agent systems. By saving the memory occupied by the replay buffer, PQN paves the way for a generation of powerful but stable algorithms that exploit end-to-end GPU vectorised deep RL.

## 7 Reproducibility Statement

All our experiments can be replicated with the repository provided anonymously here: [https://anonymous.4open.science/r/purejaxql-7283/](https://anonymous.4open.science/r/purejaxql-7283/). The repository will be publicly available after the review process. Proofs for all theorems and corollaries can be found in Appendix B.

## Appendix A Related Work

### A.1 Asynchronous Methods and Parallelisation of Q 𝑄 Q italic_Q -learning

Existing attempts to parallelise $Q$-learning adopt a distributed architecture, where a separate process continually trains the agent and sampling occurs in parallel threads which contain a delayed copy of its parameters (Horgan et al., ; Kapturowski et al., ; Badia et al., ; Hoffman et al., ). On the contrary, PQN samples and trains in the same process, enabling end-to-end single GPU training. While distributed methods can benefit from a separate process that continuously train the network, PQN is easier to implement, doesn’t introduce time-lags between the learning agent and the exploratory policy. Moreover, PQN can be optimised to be sample efficient other than fast, while distributed system usually ignore sample efficiency.

Mnih et al. () propose an asynchronous $Q$-learning, a parallelised version of $Q$-learning which performs updates of a centralised network asynchronously without using a replay buffer. Compared to PQN, asynchronous $Q$-learning still uses target networks and accumulates gradients over many timesteps to update the network. Moreover, it is a multi-threaded approach where each worker independently performs exploration and gradient updates with its own target network. This setup results in each actor being optimisd independently with its own experiences and objective, introducing significant noise into the central learner that periodically unifies the gradients. Finally, the algorithm relies on collecting historical data: "We also accumulate gradients over multiple timesteps before they are applied" (Mnih et al., ). This undermines a key benefit of parallelised methods, which is avoiding the use of data collected under historic policies (see Section 4.1).

PQN is a synchronous method where a single actor interacts with vectorised environments, and a single gradient is computed at once using all the experiences. PQN can be seen as the synchronous version of that asynchronous Q-Learning algorithm, which was never implemented before. Note that moving from asynchronous to synchronous, removing the target networks, and avoiding multi-step gradient accumulation drastically changes the optimisation procedure and implementation, resulting in a much simpler and more stable algorithm. To our knowledge, we are the first to unlock the potential of a parallelised deep $Q$-learning algorithm without a replay buffer and target networks.

### A.2 Regularisation in RL

CrossQ is a recently developed off-policy algorithm based on SAC that removes target networks and introduces BatchNorm into the critic (Bhatt et al., ). CrossQ demonstrates impressive performance when evaluated across a suite of MuJoCo domain with high sample and computational efficiency, however no analysis is provided explaining its stability properties or what effect introducing BatchNorm has. To develop PQN, we perform a rigorous analysis of both BatchNorm and LayerNorm in TD. Moreover, *CrossQ still relies on a replay buffer*. In contrast, PQN does not require storing a replay buffer of updates, allowing us to exploit truly asynchronous updates in parallelised environments using highly off-policy data.

The benefits of regularisation has also been reported in other areas of the RL literature.
Lyle et al. (Lyle et al., ; ) investigate plasticity loss in off-policy RL, a
phenomenon where neural networks lose their ability to fit a new target functions throughout training. They propose LayerNorm (Lyle et al., ) and LayerNorm with $\ell_{2}$ regularisation (Lyle et al., ), as a solution to this problem, and show improved performance on the Atari Learning Environment, but they also use other methods of stabilisation, such as target networks, that we explicitly remove. In addition, they provide no formal analysis explaining stability.

### A.3 Analysis of TD

Most prior approaches analysing TD focus on linear function approximation. Tsitsiklis & Van Roy () first proved convergence of linear, on-policy TD, arguing that the projected Bellman operator in this setting is a contraction. Dalal et al. () give the first finite time bounds for linear TD(0), under an i.i.d. data model similar to the one that we use here. Bhandari et al. () provide bounds for linear TD in both the i.i.d. and Markov chain setting. Srikant & Ying () approach the problem from the perspective of Ordinary Differential Equations (ODE) analysis, bounding the divergence of a Lyapunov function from the limiting point of the ODE that arises from the TD update scheme. Analysis of pure TD in the general nonlinear and Markov chain sampling regime is lacking.

Two papers that are most closely related to our work are: (Fellows et al., ) and (Yue et al., ).

(Yue et al., ) analyses the effect of LayerNorm in TD, however there are several important differences. Firstly, the paper analyzes the neural tangent kernel (NTK) of the update, which only exists in the limit of infinite width networks and does not capture the nonlinear instability that we analyse. We make no such assumption as this will never hold in practice. Instead, we use the analysis of Fellows et al. () which predates (Yue et al., ) and provides a more general framework for studying TD with finite width *nonlinear networks*. Moreover, Fellows et al. () provide key results on establishing stability of general TD using an eigenvalue analysis that are remarkably similar to Yue et al. ()’s SEEM framework (in particular, their key result Theorem 3 predicts divergence when
is greater than 0 is exactly the same (but less general - it’s in a linear case) as the condition in Fellows et al. () that
has strictly negative eigenvalues). We extend these results to Markov chain sampling with normalised regularisation.

Yue et al. () claim that LayerNorm alone can stabilise TD. Under our more general and applicable analysis, as our results show, LayerNorm without $\ell_{2}$ regularisation *cannot completely stabilise TD for all domains*. This is because for stability, the Jacobian eigenvalues need to be strictly negative. As Lemma 1 shows, there may still be a residual positive term that prevents this. Our empirical results in Baird’s counterexample confirm this, showing that the algorithm can only be stabilised using normalisation. Existing empirical research (Lyle et al., ) also supports this.

### A.4 Regularisation in RL

CrossQ is a recently developed off-policy algorithm based on SAC that removes target networks and introduces BatchNorm into the critic (Bhatt et al., ). CrossQ demonstrates impressive performance when evaluated across a suite of MuJoCo domain with high sample and computational efficiency, however no analysis is provided explaining its stability properties or what effect introducing BatchNorm has. To develop PQN, we perform a rigorous analysis of both BatchNorm and LayerNorm in TD. Here is a complete list of the differences between CrossQ and PQN:

- •
CrossQ is based on a soft-actor critic architecture for continuous action control. Its entropy-based actor objective optimises a stochastic policy. On the contrary, PQN consists of a single, simple value network optimised with the standard Bellman Equations, which is used to learn a deterministic policy for discrete actions.
- •
CrossQ is not parallelised, i.e., it interacts with a single environment at a time, while a fundamental contribution of PQN is handling parallel environments for faster training on modern hardware. Parallelisation of Q-Learning algorithms is not trivial: one cannot simply interact with multiple environments while leaving the rest of the learning pipeline unchanged, as this drastically modifies the ratio between interactions with the environments and the number of gradient updates. PQN approaches this problem by offering a sample-efficient implementation based on normalisation, Q-Lambda, and mini-batches/mini-epochs updates.
- •
CrossQ uses a replay buffer to perform updates, while PQN replaces the replay buffer by obtaining mini-batches directly from interactions with parallel environments.
- •
CrossQ is not directly compatible with Q-Lambda and Recurrent Neural Networks because of the overhead introduced by the replay buffers: the use of old experiences in the update step makes computation of Q-Lambda unsafe, and the use of hidden states for the RNNs problematic. To include these methods in CrossQ one should add, e.g., Retrace and Burn-In techniques. Conversely, the absence of the replay buffer in PQN allows us to use them out of the box. Note that these are crucial in many scenarios (see Q-Lambda ablation for Atari and MLP-RNN results in Craftax).
- •
The analysis of normalisation in CrossQ is purely empirical and limited to six Mujoco continuous-action tasks. This is not sufficient to make any reasonable claims for its performance in general RL scenarios. We give a theoretical basis for our method and we compare it 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). The limited evaluation provided for CrossQ is concerning, and the results in Mujoco might not reflect its true capabilities. Our results in Atari demonstrate it.
- •
PQN is designed to run fully on GPU and be compatible with end-to-end compilation, which is a fundamental step for bringing Q-Learning into modern RL research (currently dominated by PPO). CrossQ does not tackle this problem at all and offers a very standard pipeline (which consists of interacting with one environment - sampling from the replay buffer - updating the network - repeat) with the addition of normalisation. To be clear, this pipeline is exactly the same as that used by DQN in our Atari experiments, where we show that PQN is between 50x and 100x faster and uses 26 times less memory.

### A.5 Multi-Step Q 𝑄 Q italic_Q -learning

The concept of n-step returns in reinforcement learning extends the traditional one-step update to consider rewards over multiple future steps. The $n$-step return for a state-action pair (s, a) is defined as the cumulative reward over the next n steps plus the discounted value of the state reached after n steps. Several variations of n-step $Q$-learning have been proposed to enhance learning efficiency and stability. Peng & Williams () introduced a variation known as Q($\lambda$), which integrates eligibility traces to account for multiple time steps while maintaining the off-policy nature of $Q$-learning. Replay buffers are difficult to combine with $Q(\lambda)$, so standard methods like DQN use a single-step TD learning. The most relevant work that aimed to use $Q(\lambda)$ with a replay buffer is Retrace (Munos et al., ). More recent methods have tried to reconcile $\lambda$-returns with the experience buffer Daley & Amato () most notably in TD3 Kozuno et al. ().

### A.6 Multi-Agent Deep Q Learning

$Q$-learning methods are a popular choice for multi-agent RL (MARL), especially in the purely cooperative centralised training with decentralised execution (CTDE) setting (Foerster et al., ; Lowe et al., ). In CTDE, global information is made available at training time, but not at test time.
Many of these methods develop approaches to combine individual utility functions into a joint estimate of the Q function: Son et al. () introduce the individual-global-max (IGM) principle to describe when a centralised Q function can be computed from individual utility functions in a decentralised fashion; Value Decomposition Networks (VDN)  (Sunehag et al., ) combines individual value estimates by summing them, and QMIX (Rashid et al., ) learns a hypernetwork with positive weights to ensure monotonicity. All these methods can be combined with PQN, which parallises the learning process.

IPPO (De Witt et al., ) and MAPPO (Yu et al., ) use vectorised environments, adapting a single-agent method for use in multi-agent RL. These are both on-policy actor-critic based methods based on PPO.

## Appendix B Proofs

### B.1 Theorem 1 - Analysing TD

###### Theorem 1 (TD Stability) .

Let Assumptions 1 and 2 hold. If the TD criterion holds then the TD updates in Eq. 3 converge with:

$$ $\displaystyle\lim_{i\rightarrow\infty}\mathbb{E}\left[\left\lVert\phi_{i}-\phi ^{\star}\right\rVert^{2}\right]=0.$ (13) $$

###### Proof.

We use the notation $\mathbb{E}_{-\varsigma_{i}}[\cdot]$ to denote the expectation over $\{\varsigma_{0},\dots\varsigma_{i-1}\}$ and $\mathbb{E}_{\varsigma_{i}|\varsigma_{i-1}}[\cdot]$ to denote the expectation over $\varsigma_{i}$ conditioned on $\varsigma_{i-1}$.
Substituting for $\phi_{i+1}=\phi_{i}+\alpha_{i}\delta(\phi_{i},\varsigma_{i})$ into $\mathbb{E}\left[\left\lVert\phi_{i+1}-\phi^{\star}\right\rVert^{2}\right]$ yields:

$$ $\displaystyle\mathbb{E}_{\varsigma_{i},-\varsigma_{i}}$ $\displaystyle\left[\left\lVert\phi_{i+1}-\phi^{\star}\right\rVert^{2}\right]= \mathbb{E}_{\varsigma_{i},-\varsigma_{i}}\left[\left\lVert\phi_{i}+\alpha_{i} \delta(\phi_{i},\varsigma_{i})-\phi^{\star}\right\rVert^{2}\right],$ (14) $\displaystyle=$ $\displaystyle\mathbb{E}_{\varsigma_{i},-\varsigma_{i}}\left[\left\lVert\phi_{i }-\phi^{\star}\right\rVert^{2}+2\alpha_{i}\delta(\phi_{i},\varsigma_{i})^{\top }(\phi_{i}-\phi^{\star})+\alpha_{i}^{2}\lVert\delta(\phi_{i},\varsigma_{i}) \rVert^{2}\right],$ (15) $\displaystyle=$ $\displaystyle\mathbb{E}_{-\varsigma_{i}}\left[\left\lVert\phi_{i}-\phi^{\star} \right\rVert^{2}+2\alpha_{i}\mathbb{E}_{\varsigma_{i}|-\varsigma_{i}}\left[ \delta(\phi_{i},\varsigma_{i})\right]^{\top}(\phi_{i}-\phi^{\star})+\alpha_{i} ^{2}\mathbb{E}_{\varsigma_{i}|-\varsigma_{i}}\left[\lVert\delta(\phi_{i}, \varsigma_{i})\rVert^{2}\right]\right],$ (16) $\displaystyle=$ $\displaystyle\mathbb{E}_{-\varsigma_{i}}\Big{[}\left\lVert\phi_{i}-\phi^{\star }\right\rVert^{2}+2\alpha_{i}\left(\mathbb{E}_{\varsigma_{i}|-\varsigma_{i}} \left[\delta(\phi_{i},\varsigma_{i})\right]-\delta(\phi_{i})+\delta(\phi_{i}) \right)^{\top}(\phi_{i}-\phi^{\star})$ (17) $\displaystyle+\alpha_{i}^{2}\mathbb{E}_{\varsigma_{i}|-\varsigma_{i}}\left[ \lVert\delta(\phi_{i},\varsigma_{i})\rVert^{2}\right]\Big{]},$ (18) $\displaystyle=$ $\displaystyle\mathbb{E}_{-\varsigma_{i}}\Big{[}\left\lVert\phi_{i}-\phi^{\star }\right\rVert^{2}+2\alpha_{i}\delta(\phi_{i})^{\top}(\phi_{i}-\phi^{\star})$ (19) $\displaystyle+2\alpha_{i}\left(\mathbb{E}_{\varsigma_{i}|-\varsigma_{i}}\left[ \delta(\phi_{i},\varsigma_{i})\right]-\delta(\phi_{i})\right)^{\top}(\phi_{i}- \phi^{\star})+\alpha_{i}^{2}\mathbb{E}_{\varsigma_{i}|-\varsigma_{i}}\left[ \lVert\delta(\phi_{i},\varsigma_{i})\rVert^{2}\right]\Big{]},$ (20) $\displaystyle\leq$ $\displaystyle\mathbb{E}_{-\varsigma_{i}}\left[\left\lVert\phi_{i}-\phi^{\star} \right\rVert^{2}+2\alpha_{i}\delta(\phi_{i})^{\top}(\phi_{i}-\phi^{\star})\right]$ (21) $\displaystyle+2\alpha_{i}\left\lvert\mathbb{E}_{-\varsigma_{i}}\left[\left( \mathbb{E}_{\varsigma_{i}|-\varsigma_{i}}\left[\delta(\phi_{i},\varsigma_{i}) \right]-\delta(\phi_{i})\right)^{\top}\left(\phi_{i}-\phi^{\star}\right)\right ]\right\rvert+\alpha_{i}^{2}\mathbb{E}_{\varsigma_{i},-\varsigma_{i}}\left[ \lVert\delta(\phi_{i},\varsigma_{i})\rVert^{2}\right],$ (22) $$

where we have isolated the contribution of variance and non-i.i.d. sampling in deriving the final line. We now bound the non-i.i.d. contribution in total variation and variance term using Lemma 1:

$$ $\displaystyle\mathbb{E}_{\varsigma_{i},-\varsigma_{i}}\left[\left\lVert\phi_{i +1}-\phi^{\star}\right\rVert^{2}\right]\leq$ $\displaystyle\mathbb{E}_{-\varsigma_{i}}\left[\left\lVert\phi_{i}-\phi^{\star} \right\rVert^{2}+2\alpha_{i}\delta(\phi_{i})^{\top}(\phi_{i}-\phi^{\star}) \right]+2\alpha_{i}C_{\textrm{Markov}}\rho^{i}+\alpha_{i}^{2}C_{\textrm{Var}}.$ (23) $$

Note that for i.i.d. sampling, $\mathbb{E}_{\varsigma_{i}|-\varsigma_{i}}\left[\delta(\phi_{i},\varsigma_{i})
\right]=\mathbb{E}_{\varsigma_{i}\sim P_{\varsigma}}\left[\delta(\phi_{i},
\varsigma_{i})\right]=\delta(\phi_{i})$ and so $C_{\textrm{Markov}}=0$. Next, we re-write $\delta(\phi_{i})$ to contain a factor of $\phi_{i}-\phi^{\star}$. Define the line joining $\phi^{\star}$ to $\phi_{i}$ as $\ell(l)=\phi_{i}-l(\phi_{i}-\phi^{\star})$. Under 2, we can apply the fundamental theorem of calculus to integrate along this line, yielding:

$$ $\displaystyle\delta(\phi_{i})$ $\displaystyle=\delta(\phi_{i})-\underbrace{\delta(\phi^{\star})}_{=0},$ (24) $\displaystyle=\delta(\phi=\ell(0))-\delta(\phi=\ell(1)),$ (25) $\displaystyle=-\int_{0}^{1}\partial_{l}\delta(\phi=\ell(l))dl,$ (26) $\displaystyle=\int_{0}^{1}\nabla_{\phi}\delta(\phi=\ell(l))(\phi_{i}-\phi^{ \star})dl,$ (27) $\displaystyle=\int_{0}^{1}J(\phi=\ell(l))dl(\phi_{i}-\phi^{\star}),$ (28) $$

where we have used the chain rule to derive the penultimate line. Substituting yields:

$$ $\displaystyle\mathbb{E}_{\varsigma_{i},-\varsigma_{i}}\left[\left\lVert\phi_{i +1}-\phi^{\star}\right\rVert^{2}\right]\leq$ $\displaystyle\mathbb{E}_{-\varsigma_{i}}\left[\left\lVert\phi_{i}-\phi^{\star} \right\rVert^{2}+2\alpha_{i}(\phi_{i}-\phi^{\star})^{\top}\int_{0}^{1}J(\phi= \ell(l))dl(\phi_{i}-\phi^{\star})\right]$ (29) $\displaystyle+2\alpha_{i}C_{\textrm{Markov}}\rho^{i}+\alpha_{i}^{2}C_{\textrm{ Var}}.$ (30) $$

Now, as the TD criterion: $v^{\top}J(\phi)v<0$ holds almost everywhere, it follows that:

$$ $\displaystyle(\phi_{i}-\phi^{\star})^{\top}\int_{0}^{1}J(\phi=\ell(l))dl(\phi_ {i}-\phi^{\star})$ $\displaystyle<0,$ (31) $\displaystyle\implies(\phi_{i}-\phi^{\star})^{\top}\int_{0}^{1}J(\phi=\ell(l)) dl(\phi_{i}-\phi^{\star})$ $\displaystyle\leq-\lambda_{\textrm{min}}\lVert\phi_{i}-\phi^{\star}\rVert^{2},$ (32) $$

where $-\lambda_{\textrm{min}}<0$ is the smallest (in magnitude) eigenvalue of $\int_{0}^{1}J(\phi=\ell(l))dl$. Substituting yields:

$$ $\displaystyle\mathbb{E}_{\varsigma_{i},-\varsigma_{i}}$ $\displaystyle\left[\left\lVert\phi_{i+1}-\phi^{\star}\right\rVert^{2}\right] \leq\mathbb{E}_{-\varsigma_{i}}\left[\left\lVert\phi_{i}-\phi^{\star}\right \rVert^{2}\right](1-2\alpha_{i}\lambda_{\textrm{min}})+2\alpha_{i}C_{\textrm{ Markov}}\rho^{i}+\alpha_{i}^{2}C_{\textrm{Var}}.$ (33) $$

Re-arranging yields:

$$ $\displaystyle 2\lambda_{\textrm{min}}\alpha_{i}\mathbb{E}_{-\varsigma_{i}}$ $\displaystyle\left[\left\lVert\phi_{i}-\phi^{\star}\right\rVert^{2}\right]$ (34) $\displaystyle\leq\mathbb{E}_{-\varsigma_{i}}\left[\left\lVert\phi_{i}-\phi^{ \star}\right\rVert^{2}\right]-\mathbb{E}_{\varsigma_{i},-\varsigma_{i}}\left[ \left\lVert\phi_{i+1}-\phi^{\star}\right\rVert^{2}\right]+2\alpha_{i}C_{ \textrm{Markov}}\rho^{i}+\alpha_{i}^{2}C_{\textrm{Var}}.$ (35) $$

Summing over $i$ up to timestep $t$ yields:

$$ $\displaystyle 2\lambda_{\textrm{min}}$ $\displaystyle\sum_{i=0}^{t}\alpha_{i}\mathbb{E}_{-\varsigma_{i}}\left[\left \lVert\phi_{i}-\phi^{\star}\right\rVert^{2}\right]$ (36) $\displaystyle\leq\mathbb{E}_{-\varsigma_{i}}\left[\left\lVert\phi_{0}-\phi^{ \star}\right\rVert^{2}\right]-\mathbb{E}_{\varsigma_{i},-\varsigma_{i}}\left[ \left\lVert\phi_{t+1}-\phi^{\star}\right\rVert^{2}\right]+2C_{\textrm{Markov}} \sum_{i=0}^{t}\alpha_{i}\rho^{i}+C_{\textrm{Var}}\sum_{i=0}^{t}\alpha_{i}^{2},$ (37) $\displaystyle\leq\mathbb{E}_{-\varsigma_{i}}\left[\left\lVert\phi_{0}-\phi^{ \star}\right\rVert^{2}\right]+2C_{\textrm{Markov}}\sum_{i=0}^{t}\alpha_{i}\rho ^{i}+C_{\textrm{Var}}\sum_{i=0}^{t}\alpha_{i}^{2},$ (38) $\displaystyle\implies$ $\displaystyle\sum_{i=0}^{t}\frac{\alpha_{i}}{\sum_{i=0}^{t}\alpha_{i}}\mathbb{ E}_{-\varsigma_{i}}\left[\left\lVert\phi_{i}-\phi^{\star}\right\rVert^{2}\right]$ (39) $\displaystyle\leq\frac{1}{2\lambda_{\textrm{min}}}\left(\frac{\mathbb{E}_{- \varsigma_{i}}\left[\left\lVert\phi_{0}-\phi^{\star}\right\rVert^{2}\right]}{ \sum_{i=0}^{t}\alpha_{i}}+2C_{\textrm{Markov}}\frac{\sum_{i=0}^{t}\alpha_{i} \rho^{i}}{\sum_{i=0}^{t}\alpha_{i}}+C_{\textrm{Var}}\frac{\sum_{i=0}^{t}\alpha _{i}^{2}}{\sum_{i=0}^{t}\alpha_{i}}\right),$ (40) $$

where the penultimate bound follows from $\mathbb{E}_{\varsigma_{i},-\varsigma_{i}}\left[\left\lVert\phi_{t+1}-\phi^{
\star}\right\rVert^{2}\right]\geq 0$. In preparation for taking the limit $t\rightarrow\infty$, we observe that by the Cauchy-Schwartz inequality:

$$ $\displaystyle\sum_{i=0}^{t}\alpha_{i}\rho^{i}=\sum_{i=0}^{t}\lvert\alpha_{i} \rvert\lvert\rho^{i}\rvert\leq\sqrt{\sum_{i=0}^{t}\alpha_{i}^{2}\rho^{2}i}\leq \sqrt{\sum_{i=0}^{t}\alpha_{i}^{2}\sum_{i=0}^{t}\rho^{2i}}$ (41) $$

Now, from 1, $\lim_{t\rightarrow\infty}\sum_{i=0}^{t}\alpha_{i}^{2}<\infty$, hence:

$$ $\displaystyle\lim_{t\rightarrow\infty}\sum_{i=0}^{t}\alpha_{i}\rho^{i}\leq \sqrt{\lim_{t\rightarrow\infty}\sum_{i=0}^{t}\alpha_{i}^{2}\lim_{t\rightarrow \infty}\sum_{i=0}^{t}\rho^{2i}}=\mathcal{O}(1)$ (42) $$

As $\lim_{t\rightarrow\infty}\sum_{i=0}^{t}\alpha_{i}=\infty$, this implies:

$$ $\displaystyle\lim_{t\rightarrow\infty}\frac{\sum_{i=0}^{t}\alpha_{i}\rho^{i}}{ \sum_{i=0}^{t}\alpha_{i}}=0.$ (43) $$

We are now ready to take limits of Inq. 40, yielding:

$$ $\displaystyle\lim_{t\rightarrow\infty}\sum_{i=0}^{t}\frac{\alpha_{i}}{\sum_{i= 0}^{t}\alpha_{i}}\mathbb{E}_{-\varsigma_{i}}\left[\left\lVert\phi_{i}-\phi^{ \star}\right\rVert^{2}\right]=0.$ (44) $$

Eq. 44 proves our desired result:

$$ $\displaystyle\lim_{i\rightarrow\infty}\mathbb{E}_{-\varsigma_{i}}\left[\left \lVert\phi_{i}-\phi^{\star}\right\rVert^{2}\right]=0.$ (45) $$

To see why, assume this does not hold, that is $\lim_{i\rightarrow\infty}\mathbb{E}_{-\varsigma_{i}}\left[\left\lVert\phi_{i}-
\phi^{\star}\right\rVert^{2}\right]\neq 0$. This implies there exists some infinite length sub-sequence $S$ such that for all $i\in S$:

$$ $\displaystyle\mathbb{E}_{-\varsigma_{i}}\left[\left\lVert\phi_{i}-\phi^{\star} \right\rVert^{2}\right]>0,$ (46) $$

hence, as all quantities are positive:

$$ $\displaystyle\lim_{t\rightarrow\infty}\sum_{i=0}^{t}\frac{\alpha_{i}}{\sum_{i= 0}^{t}\alpha_{i}}\mathbb{E}_{-\varsigma_{i}}\left[\left\lVert\phi_{i}-\phi^{ \star}\right\rVert^{2}\right]\geq\lim_{t\rightarrow\infty}\sum_{i\in S}\frac{ \alpha_{i}}{\sum_{i=0}^{t}\alpha_{i}}\mathbb{E}_{-\varsigma_{i}}\left[\left \lVert\phi_{i}-\phi^{\star}\right\rVert^{2}\right]>0,$ (47) $$

which is a contradiction.
∎

###### Lemma 1 .

Let 2 hold. Then there exist constants: $0<C_{\textnormal{{Markov}}}<\infty$, $0<C_{\textnormal{{Var}}}<\infty$ and $\rho\in[0,1)$ such that:

$$ $\displaystyle\left\lvert\mathbb{E}_{-\varsigma_{i}}\left[\left(\mathbb{E}_{ \varsigma_{i}|-\varsigma_{i}}\left[\delta(\phi_{i},\varsigma_{i})\right]- \delta(\phi_{i})\right)^{\top}\left(\phi_{i}-\phi^{\star}\right)\right]\right \rvert\leq C_{\textrm{Markov}}\rho^{i},\quad\mathbb{E}_{\varsigma_{i},- \varsigma_{i}}\left[\lVert\delta(\phi_{i},\varsigma_{i})\rVert^{2}\right]\leq C _{\textnormal{{Var}}}.$ (48) $$

###### Proof.

For both results, we use the fact that, because $\Phi$ is compact and $\mathcal{X}$ is bounded, rewards are bounded and $\delta(\phi,\varsigma)$ is Lipschitz under 2, $\delta(\phi,\varsigma)$ is bounded almost everywhere.

To prove the first bound, we denote the marginal probability distribution of the $i$-th element in the Markov chain $\varsigma_{i}$ as $P^{i}$ with density:

$$ $\displaystyle p^{i}(\varsigma_{i})=\int p(-\varsigma_{i},\varsigma_{i})d(- \varsigma_{i}).$ (49) $$

Under this notation we write:

$$ $\displaystyle\mathbb{E}_{-\varsigma_{i}}$ $\displaystyle\left[\left(\mathbb{E}_{\varsigma_{i}|-\varsigma_{i}}\left[\delta (\phi_{i},\varsigma_{i})\right]-\delta(\phi_{i})\right)^{\top}\left(\phi_{i}- \phi^{\star}\right)\right]$ (50) $\displaystyle=$ $\displaystyle\mathbb{E}_{-\varsigma_{i},\varsigma_{i}}\left[\left(\delta(\phi_ {i},\varsigma_{i})-\mathbb{E}_{\varsigma_{i}^{\prime}\sim P_{\varsigma}}\left[ \delta(\phi_{i},\varsigma_{i}^{\prime})\right]\right)^{\top}\left(\phi_{i}- \phi^{\star}\right)\right],$ (51) $\displaystyle=$ $\displaystyle\mathbb{E}_{-\varsigma_{i},\varsigma_{i}}\left[\delta(\phi_{i}, \varsigma_{i})^{\top}\left(\phi_{i}-\phi^{\star}\right)-\mathbb{E}_{\varsigma_ {i}^{\prime}\sim P_{\varsigma}}\left[\delta(\phi_{i},\varsigma_{i}^{\prime}) \right]^{\top}\left(\phi_{i}-\phi^{\star}\right)\right],$ (52) $\displaystyle=$ $\displaystyle\mathbb{E}_{-\varsigma_{i},\varsigma_{i}}\left[\delta(\phi_{i}, \varsigma_{i})^{\top}\left(\phi_{i}-\phi^{\star}\right)-\mathbb{E}_{\varsigma_ {i}^{\prime}\sim P_{\varsigma}}\left[\delta(\phi_{i},\varsigma_{i}^{\prime})^{ \top}\left(\phi_{i}-\phi^{\star}\right)\right]\right],$ (53) $\displaystyle=$ $\displaystyle\mathbb{E}_{\varsigma_{i}\sim P^{i}}\left[\mathbb{E}_{-\varsigma_ {i}\sim P^{i}(\varsigma_{i})}\left[\delta(\phi_{i},\varsigma_{i})^{\top}\left( \phi_{i}-\phi^{\star}\right)\right]\right]-\mathbb{E}_{\varsigma_{i}\sim P_{ \varsigma}}\left[\mathbb{E}_{-\varsigma_{i}\sim P^{i}(\varsigma_{i})}\left[ \delta(\phi_{i},\varsigma_{i})^{\top}\left(\phi_{i}-\phi^{\star}\right)\right] \right],$ (54) $$

where $P^{i}(\varsigma_{i})$ is the backwards conditional distribution in the Markov chain with density:

$$ $\displaystyle p^{i}(-\varsigma_{i}|\varsigma_{i})=\frac{p^{i}(-\varsigma_{i}, \varsigma_{i})}{p^{i}(\varsigma_{i})}.$ (55) $$

Introducing the notation:

$$ $\displaystyle g(\varsigma_{i})=\mathbb{E}_{-\varsigma_{i}\sim P^{i}(\varsigma_ {i})}\left[\delta(\phi_{i},\varsigma_{i})^{\top}\left(\phi_{i}-\phi^{\star} \right)\right],$ (56) $$

we write Eq. 54 as:

$$ $\displaystyle\mathbb{E}_{-\varsigma_{i}}$ $\displaystyle\left[\left(\mathbb{E}_{\varsigma_{i}|-\varsigma_{i}}\left[\delta (\phi_{i},\varsigma_{i})\right]-\delta(\phi_{i})\right)^{\top}\left(\phi_{i}- \phi^{\star}\right)\right]=\mathbb{E}_{\varsigma_{i}\sim P^{i}}\left[g( \varsigma_{i})\right]-\mathbb{E}_{\varsigma_{i}\sim P_{\varsigma}}\left[g( \varsigma_{i})\right],$ (57) $\displaystyle=\mathbb{E}_{\varsigma_{0}}\left[\mathbb{E}_{\varsigma_{i}\sim P^ {i}(\varsigma_{0})}\left[g(\varsigma_{i})\right]-\mathbb{E}_{\varsigma_{i}\sim P _{\varsigma}}\left[g(\varsigma_{i})\right]\right],$ (58) $\displaystyle=\mathbb{E}_{\varsigma_{0}}\left[g_{\textrm{max}}\left(\mathbb{E} _{\varsigma_{i}\sim P^{i}(\varsigma_{0})}\left[\frac{g(\varsigma_{i})}{g_{ \textrm{max}}}\right]-\mathbb{E}_{\varsigma_{i}\sim P_{\varsigma}}\left[\frac{ g(\varsigma_{i})}{g_{\textrm{max}}}\right]\right)\right],$ (59) $$

where $g_{\textrm{max}}\coloneqq\max_{\varsigma}\lvert g(\varsigma)\rvert<\infty$ almost everywhere, which follows from the fact that $\delta(\phi,\varsigma)$ is bounded almost everywhere, implying $g(\varsigma)$ is also bounded almost everywhere.
Now, as $\frac{g(\cdot)}{g_{\textrm{max}}}:\mathcal{X}\times\mathbb{R}\times\mathcal{X}
\rightarrow[-1,1]$, we can bound Eq. 59 in total variation using Roberts & Rosenthal ()[Proposition 3b]:

$$ $\displaystyle\bigg{\lvert}\mathbb{E}_{\varsigma_{0}}$ $\displaystyle\left[g_{\textrm{max}}\left(\mathbb{E}_{\varsigma_{i}\sim P^{i}( \varsigma_{0})}\left[\frac{g(\varsigma_{i})}{g_{\textrm{max}}}\right]-\mathbb{ E}_{\varsigma_{i}\sim P_{\varsigma}}\left[\frac{g(\varsigma_{i})}{g_{\textrm{ max}}}\right]\right)\right]\bigg{\rvert}\leq$ (60) $\displaystyle\mathbb{E}_{\varsigma_{0}}\left[\left\lvert g_{\textrm{max}}\left (\mathbb{E}_{\varsigma_{i}\sim P^{i}(\varsigma_{0})}\left[\frac{g(\varsigma_{i })}{g_{\textrm{max}}}\right]-\mathbb{E}_{\varsigma_{i}\sim P_{\varsigma}}\left [\frac{g(\varsigma_{i})}{g_{\textrm{max}}}\right]\right)\right\rvert\right]$ (61) $\displaystyle=$ $\displaystyle 2g_{\textrm{max}}\mathbb{E}_{\varsigma_{0}}\left[\frac{1}{2} \left\lvert\mathbb{E}_{\varsigma_{i}\sim P^{i}(\varsigma_{0})}\left[\frac{g( \varsigma_{i})}{g_{\textrm{max}}}\right]-\mathbb{E}_{\varsigma_{i}\sim P_{ \varsigma}}\left[\frac{g(\varsigma_{i})}{g_{\textrm{max}}}\right]\right\rvert \right],$ (62) $\displaystyle\leq$ $\displaystyle 2g_{\textrm{max}}\mathbb{E}_{\varsigma_{0}}\left[\frac{1}{2}\sup _{f:\mathcal{X}\times\mathbb{R}\times\mathcal{X}\rightarrow[-1,1]}\left\lvert \mathbb{E}_{\varsigma_{i}\sim P^{i}(\varsigma_{0})}\left[f(\varsigma_{i}) \right]-\mathbb{E}_{\varsigma_{i}\sim P_{\varsigma}}\left[f(\varsigma_{i}) \right]\right\rvert\right],$ (63) $\displaystyle=$ $\displaystyle 2g_{\textrm{max}}\mathbb{E}_{\varsigma_{0}}\left[\textrm{TV}(P^{ i}(\varsigma_{0})\|P_{\varsigma})\right],$ (64) $$

where $\textrm{TV}(P^{i}(\varsigma_{0})\|P_{\varsigma})$ is the total variational distance between the marginal distribution $P^{i}(\varsigma_{0})$ (conditioned on initial observations) and the steady state distribution $P_{\varsigma}$. Now, as the Markov chain is geometricaly ergodic, by definition there exists some function $M(\varsigma_{0})$ and constant $\rho\in[0,1)$ such that:

$$ $\displaystyle\textrm{TV}(P^{i}(\varsigma_{0})\|P_{\varsigma})\leq M(\varsigma_ {0})\rho^{i},$ (65) $$

almost surely (see Roberts & Rosenthal ()[Section 3.4]), hence substituting into Eq. 64 yields our desired result:

$$ $\displaystyle\left\lvert\mathbb{E}_{-\varsigma_{i}}\left[\left(\mathbb{E}_{ \varsigma_{i}|-\varsigma_{i}}\left[\delta(\phi_{i},\varsigma_{i})\right]- \delta(\phi_{i})\right)^{\top}\left(\phi_{i}-\phi^{\star}\right)\right]\right\rvert\leq$ $\displaystyle 2g_{\textrm{max}}\mathbb{E}_{\varsigma_{0}}\left[\textrm{TV}(P^{ i}(\varsigma_{0})\|P_{\varsigma})\right],$ (66) $\displaystyle\leq$ $\displaystyle 2g_{\textrm{max}}\mathbb{E}_{\varsigma_{0}}\left[M(\varsigma_{0} )\rho^{i}\right],$ (67) $\displaystyle=$ $\displaystyle 2g_{\textrm{max}}\mathbb{E}_{\varsigma_{0}}\left[M(\varsigma_{0} )\right]\rho^{i},$ (68) $\displaystyle=$ $\displaystyle C_{\textrm{Markov}}\rho^{i},$ (69) $$

where $C_{\textrm{Markov}}\coloneqq 2g_{\textrm{max}}\mathbb{E}_{\varsigma_{0}}\left[
M(\varsigma_{0})\right]<\infty$.

Our second bound follows from the fact that $\delta(\phi,\varsigma)$ is bounded almost everywhere. This implies there exists some $C_{\textnormal{{Var}}}>0$ such that $\lVert\delta(\phi,\varsigma)\rVert\leq C_{\textnormal{{Var}}}$ almost everywhere, hence:

$$ $\displaystyle\mathbb{E}_{\varsigma_{i},-\varsigma_{i}}\left[\lVert\delta(\phi_ {i},\varsigma_{i})\rVert^{2}\right]\leq C_{\textnormal{{Var}}}.$ (70) $$

∎

###### Theorem 1 (TD Stability) .

###### Proof.

###### Lemma 1 .

###### Proof.

### B.2 Theorem 2 - Stabilising TD with LayerNorm

For all proofs, we introduce the following simplifying notation:

$$ $\displaystyle\textrm{LayerNorm}_{i}[f](x)\coloneqq$ $\displaystyle\frac{1}{\sqrt{k}}\cdot\frac{f_{i}(x)-\hat{\mu}\left[f\right](x)} {\hat{\sigma}[f](x)},$ (71) $$

where $\hat{\mu}\left[f\right](x)$ and $\hat{\sigma}[f](x)$ are the element-wise empirical mean and standard deviation of the output $f(x)$:

$$ $\displaystyle\hat{\mu}\left[f\right](x)\coloneqq\frac{1}{k}\sum_{i=0}^{k-1}f_{ i}(x),\quad\hat{\sigma}[f](x)\coloneqq\sqrt{\frac{1}{k}\sum_{i=0}^{k-1}(f_{i}( x)-\hat{\mu}\left[f\right](x))^{2}+\epsilon},$ (72) $$

We will make use of the following three key properties of LayerNorm:

###### Proposition 1 .

For any bounded $f:\mathcal{X}\rightarrow\mathbb{R}^{k}$:

$$ $\displaystyle\lVert\textnormal{{LayerNorm}}[f(x)]\rVert$ $\displaystyle\leq 1,$ (73) $\displaystyle\partial_{f_{i}}\textnormal{ {LayerNorm}}_{j}[f(x)]$ $\displaystyle=\mathcal{O}\left(k^{-\frac{1}{2}}\left(\mathbbm{1}(i=j)+\frac{1} {k}\right)\right),$ (74) $\displaystyle\partial_{f_{s}}\partial_{f_{t}}\textnormal{ {LayerNorm}}_{j}[f(x)]$ $\displaystyle=\mathcal{O}\left(k^{-\frac{3}{2}}\left(\mathbbm{1}(t=j)+\mathbbm {1}(t=s)+\mathbbm{1}(j=s)+\frac{1}{k}\right)\right).$ (75) $$

###### Proof.

Our first result follows directly from the definition of LayerNorm:

$$ $\displaystyle\lVert\textnormal{{LayerNorm}}[f(x)]\rVert$ $\displaystyle=\frac{1}{\sqrt{k}}\frac{\lVert f(x)-\hat{\mu}\left[f\right](x) \rVert}{\hat{\sigma}[f](x)},$ (76) $\displaystyle=\frac{\sqrt{\frac{1}{k}\sum_{i=0}^{k-1}\left(f_{i}(x)-\hat{\mu} \left[f\right](x)\right)^{2}}}{\hat{\sigma}[f](x)},$ (77) $\displaystyle\leq\frac{\sqrt{\frac{1}{k}\sum_{i=0}^{k-1}\left(f_{i}(x)-\hat{ \mu}\left[f\right](x)\right)^{2}+\epsilon}}{\hat{\sigma}[f](x)},$ (78) $\displaystyle=\frac{\hat{\sigma}[f](x)}{\hat{\sigma}[f](x)},$ (79) $\displaystyle=1,$ (80) $$

as required. For our second result, we take partial derivatives with respect to the $i$th input channel to the LayerNorm:

$$ $\displaystyle\partial_{f_{i}}\textnormal{ {LayerNorm}}_{j}[f(x)]$ $\displaystyle=\frac{1}{\sqrt{k}}\left(\frac{\mathbbm{1}(i=j)-\frac{1}{k}}{\hat {\sigma}[f](x)}-\frac{f_{j}(x)-\hat{\mu}[f](x)}{\hat{\sigma}[f](x)^{2}} \partial_{f_{i}}\hat{\sigma}[f](x)\right),$ (81) $\displaystyle=\frac{1}{\sqrt{k}}\left(\frac{\mathbbm{1}(i=j)-\frac{1}{k}}{\hat {\sigma}[f](x)}-\sqrt{k}\frac{\textrm{LayerNorm}_{j}[f(x)]}{\hat{\sigma}[f](x) }\partial_{f_{i}}\hat{\sigma}[f](x)\right).$ (82) $$

Finding the derivative of the empirical variance yields:

$$ $\displaystyle\partial_{f_{i}}\hat{\sigma}[f](x)$ $\displaystyle=\frac{1}{k\hat{\sigma}[f](x)}\sum_{l=0}^{k-1}(f_{l}(x)-\hat{\mu} \left[f\right](x))\left(\mathbbm{1}(i=l)-\frac{1}{k}\right),$ (83) $\displaystyle=\frac{1}{k\hat{\sigma}[f](x)}\left(\sum_{l=0}^{k-1}(f_{l}(x)- \hat{\mu}\left[f\right](x))\mathbbm{1}(i=l)-\sum_{l=0}^{k-1}(f_{l}(x)-\hat{\mu }\left[f\right](x))\frac{1}{k}\right),$ (84) $\displaystyle=\frac{f_{i}(x)-\hat{\mu}\left[f\right](x)}{k\hat{\sigma}[f](x)},$ (85) $\displaystyle=\frac{1}{\sqrt{k}}\textrm{LayerNorm}_{i}[f(x)],$ (86) $$

hence:

$$ $\displaystyle\partial_{f_{i}}$ $\displaystyle\textnormal{ {LayerNorm}}_{j}[f(x)]$ (87) $\displaystyle=\frac{1}{\sqrt{k}\hat{\sigma}[f](x)}\left(\mathbbm{1}(i=j)-\frac {1}{k}+\textrm{LayerNorm}_{i}[f(x)]\textrm{LayerNorm}_{j}[f(x)]\right),$ (88) $\displaystyle=\mathcal{O}\left(\frac{1}{\sqrt{k}}\left(\mathbbm{1}(i=j)+\frac{ 1}{k}\right)\right),$ (89) $$

where we use the fact that $\textrm{LayerNorm}_{j}[f(x)]=\mathcal{O}\left(\frac{1}{\sqrt{k}}\right)$ to derive the final line. To prove our third result, we start with the first order partial derivative using Eq. 88:

$$ $\displaystyle\partial_{f_{t}}$ $\displaystyle\textnormal{ {LayerNorm}}_{j}[f(x)]=\frac{1}{\sqrt{k}\hat{\sigma} [f](x)}\left(\mathbbm{1}(t=j)-\frac{1}{k}+\textrm{LayerNorm}_{t}[f(x)]\textrm{ LayerNorm}_{j}[f(x)]\right),$ (90) $$

Taking partial derivatives with respect to $f_{s}$ yields:

$$ $\displaystyle\partial_{f_{s}}\partial_{f_{t}}\textnormal{ {LayerNorm}}_{j}[f(x)]$ (91) $\displaystyle=-\frac{\partial_{f_{s}}\hat{\sigma}[f](x)}{\sqrt{k}\hat{\sigma}[ f](x)^{2}}\left(\mathbbm{1}(t=j)-\frac{1}{k}+\textrm{LayerNorm}_{t}[f(x)] \textrm{LayerNorm}_{j}[f(x)]\right)$ (92) $\displaystyle-k^{-\frac{1}{2}}\cdot\frac{\partial_{f_{s}}\textrm{LayerNorm}_{t }[f(x)]\cdot\textrm{LayerNorm}_{j}[f(x)]+\textrm{LayerNorm}_{t}[f(x)]\partial_ {f_{s}}\textrm{LayerNorm}_{j}[f(x)]}{\hat{\sigma}[f](x)},$ (93) $\displaystyle=\mathcal{O}\left(k^{-\frac{3}{2}}\left(\mathbbm{1}(t=j)+\frac{1} {k}\right)\right)+\mathcal{O}\left(k^{-\frac{3}{2}}\left(\mathbbm{1}(t=s)+ \frac{1}{k}\right)\right)+\mathcal{O}\left(k^{-\frac{3}{2}}\left(\mathbbm{1}(j =s)+\frac{1}{k}\right)\right),$ (94) $\displaystyle=\mathcal{O}\left(k^{-\frac{3}{2}}\left(\mathbbm{1}(t=j)+\mathbbm {1}(t=s)+\mathbbm{1}(j=s)+\frac{1}{k}\right)\right),$ (95) $$

as required.
∎

###### Theorem 2 .

Let 2 apply. Using the LayerNorm regularised TD update in Eq. 12, there exists some finite $k^{\prime}$ such that the TD stability criterion holds for all $k>k^{\prime}$

###### Proof.

From the definition of the expected regularised TD error vector:

$$ $\displaystyle\delta_{\textrm{reg}}(\phi)=$ $\displaystyle\mathbb{E}_{\varsigma\sim P_{\varsigma}}\left[\left(r+\gamma Q_{ \phi}(x^{\prime})-Q_{\phi}(x)\right)\nabla_{\phi}Q_{\phi}(x)\right]-\frac{\eta \gamma}{2}\begin{bmatrix}w\\ 0\\ \vdots\\ 0\end{bmatrix},$ (96) $\displaystyle\implies v^{\top}\nabla_{\phi}\delta_{\textrm{reg}}(\phi)$ $\displaystyle v=\mathbb{E}_{\varsigma\sim P_{\varsigma}}\left[\left(r+\gamma Q _{\phi}(x^{\prime})-Q_{\phi}(x)\right)v^{\top}\nabla_{\phi}^{2}Q_{\phi}(x)v\right]$ (97) $\displaystyle+\mathbb{E}_{\varsigma\sim P_{\varsigma}}\left[v^{\top}(\gamma \nabla_{\phi}Q_{\phi}(x^{\prime})-\nabla_{\phi}Q_{\phi}(x))\nabla_{\phi}Q_{ \phi}(x)^{\top}v\right]-\left(\frac{\eta\gamma}{2}\right)^{2}v_{w}^{\top}v_{w}.$ (98) $$

Applying Lemma 2 and taking the limit $k\rightarrow\infty$ yields:

$$ $\displaystyle\lim_{k\rightarrow\infty}v^{\top}\nabla_{\phi}\delta_{\textrm{reg }}(\phi)v=\left(\frac{\gamma}{2}\right)^{2}(1-\eta^{2})v_{w}^{\top}v_{w}<0,$ (99) $$

almost everywhere, which follows from the fact $\eta\geq 1$, hence:

$$ $\displaystyle\sup_{v,\phi}\frac{1}{v^{\top}v}v^{\top}\nabla_{\phi}\delta_{ \textrm{reg}}(\phi)v<0,$ (100) $$

almost everywhere, as required.
∎

###### Lemma 2 .

Let 2 apply. Using the LayerNorm $Q$-function defined in Eq. 8:

$$ $\displaystyle\gamma v^{\top}\nabla_{\phi}Q_{\phi}(x^{\prime})\nabla_{\phi}Q_{ \phi}(x)^{\top}v-(v^{\top}\nabla_{\phi}Q_{\phi}(x))^{2}$ $\displaystyle\leq\left(\frac{\gamma\lVert v\rVert}{2}\right)^{2}+\mathcal{O} \left(\frac{1}{\sqrt{k}}\right),$ (101) $\displaystyle\left\lvert\left(r+\gamma Q_{\phi}(x^{\prime})-Q_{\phi}(x)\right) v^{\top}\nabla_{\phi}^{2}Q_{\phi}(x)v\right\rvert$ $\displaystyle=\mathcal{O}\left(\frac{1}{\sqrt{k}}\right),$ (102) $$

almost surely for any test vector and any state-action transition pair $x,x^{\prime}\in\mathcal{X}$.

###### Proof.

We start by writing $M$ in terms of its row vectors:

$$ $\displaystyle M=\begin{bmatrix}\begin{array}[]{@{\protect\rotvert\;}c@{\; \protect\rotvert}}\rotatebox[origin={c}]{90.0}{$|$}\;\lx@intercol\hfil m_{0}^{ T}\hfil\;\rotatebox[origin={c}]{90.0}{$|$}\\ \rotatebox[origin={c}]{90.0}{$|$}\;\lx@intercol\hfil m_{1}^{\top}\hfil\; \rotatebox[origin={c}]{90.0}{$|$}\\ \lx@intercol\hfil\vdots\hfil\\ \rotatebox[origin={c}]{90.0}{$|$}\;\lx@intercol\hfil m_{k-1}^{\top}\hfil\; \rotatebox[origin={c}]{90.0}{$|$}\end{array}\end{bmatrix}.$ (103) $$

and splitting the test vector into the corresponding $k+1$ sub-vectors:

$$ $\displaystyle v^{\top}=[v_{w}^{T},v_{m_{0}}^{\top},v_{m_{1}}^{\top},\cdots v_{ m_{k-1}}^{\top}],$ (104) $$

where $v_{w}$ is a vector with the same dimension as the final weight vector $w$ and each $v_{m_{i}}\in\mathbb{R}^{n}$ has the same dimension as $x$. Let $f_{M}(x)\coloneqq\sigma\circ Mx$. Our results follow from Lemma 3 and Lemma 4.
∎

###### Lemma 3 (Mitigating Off-policy Instability) .

Under 2, using the LayerNorm critic in Eq. 8:

$$ $\displaystyle\gamma v^{\top}\nabla_{\phi}Q_{\phi}(x^{\prime})\nabla_{\phi}Q_{ \phi}(x)^{\top}v-(v^{\top}\nabla_{\phi}Q_{\phi}(x))^{2}$ $\displaystyle\leq\left(\frac{\gamma\lVert v_{w}\rVert}{2}\right)^{2}+\mathcal{ O}\left(\frac{1}{\sqrt{k}}\right),$ (105) $$

almost surely for any test vector and any state-action transition pair $x,x^{\prime}\in\mathcal{X}$.

###### Proof.

We start by split the left hand side of Eq. 105 into two terms, one determining the stability of the final layer weights and one for the matrix vectors:

$$ $\displaystyle\gamma v^{\top}\nabla_{\phi}Q_{\phi}(x^{\prime})\nabla_{\phi}Q_{ \phi}(x)^{\top}v-(v^{\top}\nabla_{\phi}Q_{\phi}(x))^{2}$ (106) $\displaystyle\quad=\gamma v_{w}^{\top}\nabla_{w}Q_{\phi}(x^{\prime})\nabla_{w} Q_{\phi}(x)^{\top}v_{w}-(v_{w}^{\top}\nabla_{w}Q_{\phi}(x))^{2}$ (107) $\displaystyle\qquad+\sum_{i=0}^{k-1}\left(\gamma v_{m_{i}}^{\top}\nabla_{m_{i} }Q_{\phi}(x^{\prime})\nabla_{m_{i}}Q_{\phi}(x)^{\top}v_{m_{i}}-v_{m_{i}}^{\top }\nabla_{m_{i}}Q_{\phi}(x)\nabla_{m_{i}}Q_{\phi}(x)^{\top}v_{m_{i}}\right).$ (108) $$

We focus on the term determining stability of the final layer weights. Taking derivatives of the critic with respect to the final layer weights $w$ yields:

$$ $\displaystyle\nabla_{w}Q_{\phi}(x)$ $\displaystyle=\textrm{LayerNorm}\left[f_{M}(x)\right].$ (109) $$

We then bound $v_{w}^{\top}\nabla_{w}Q_{\phi}(x^{\prime})\nabla_{w}Q_{\phi}(x)^{\top}v_{w}$ as:

$$ $\displaystyle v_{w}^{\top}\nabla_{w}Q_{\phi}(x^{\prime})\nabla_{w}Q_{\phi}(x)^ {\top}v_{w}$ $\displaystyle\leq\lVert v_{w}\rVert\lVert\nabla_{w}Q_{\phi}(x^{\prime})\rVert \lvert\nabla_{w}Q_{\phi}(x)^{\top}v_{w}\rvert,$ (110) $\displaystyle=\lVert v_{w}\rVert\lVert\textrm{LayerNorm}\left[f_{M}(x^{\prime} )\right]\rVert\lvert\nabla_{w}Q_{\phi}(x)^{\top}v_{w}\rvert,$ (111) $\displaystyle\leq\lVert v_{w}\rVert\lvert\nabla_{w}Q_{\phi}(x)^{\top}v_{w}\rvert,$ (112) $$

where we applied $\lVert\textrm{LayerNorm}\left[f_{M}(x^{\prime})\right]\rVert<0$ from Proposition 1 to in deriving the final line. Defining $\eta\coloneqq\lvert\nabla_{w}Q_{\phi}(x)^{\top}v_{w}\rvert$ yields:

$$ $\displaystyle\gamma v_{w}^{\top}\nabla_{w}Q_{\phi}(x^{\prime})\nabla_{w}Q_{ \phi}(x)^{\top}v_{w}-(v_{w}^{\top}\nabla_{w}Q_{\phi}(x))^{2}$ $\displaystyle\leq\gamma w_{\textrm{max}}\lvert\nabla_{w}Q_{\phi}(x)^{\top}v_{w }\rvert-\lvert\nabla_{w}Q_{\phi}(x)^{\top}v_{w}\rvert^{2},$ (113) $\displaystyle=\gamma\lVert v_{w}\rVert\eta-\eta^{2},$ (114) $\displaystyle\leq\max_{\eta}\left(\gamma\lVert v_{w}\rVert\eta-\eta^{2}\right).$ (115) $$

Our desired result follows from the fact that the function $\gamma\lVert v_{w}\rVert\eta-\eta^{2}$ is maximised at $\eta=\frac{\gamma\lVert v_{w}\rVert}{2}$, hence:

$$ $\displaystyle\gamma v_{w}^{\top}\nabla_{w}Q_{\phi}(x^{\prime})\nabla_{w}Q_{ \phi}(x)^{\top}v_{w}-(v_{w}^{\top}\nabla_{w}Q_{\phi}(x))^{2}\leq\frac{\gamma^{ 2}\lVert v_{w}\rVert}{2}-\left(\frac{\gamma\lVert v_{w}\rVert}{2}\right)^{2}$ $\displaystyle=\left(\frac{\gamma\lVert v_{w}\rVert}{2}\right)^{2}$ (116) $$

Substituting into Eq. 108 yields:

$$ $\displaystyle\gamma v^{\top}\nabla_{\phi}Q_{\phi}(x^{\prime})\nabla_{\phi}Q_{ \phi}(x)^{\top}v-(v^{\top}\nabla_{\phi}Q_{\phi}(x))^{2}\leq\left(\frac{\gamma \lVert v_{w}\rVert}{2}\right)^{2}$ (117) $\displaystyle\qquad+\sum_{i=0}^{k-1}\left(\gamma v_{m_{i}}^{\top}\nabla_{m_{i} }Q_{\phi}(x^{\prime})\nabla_{m_{i}}Q_{\phi}(x)^{\top}v_{m_{i}}-v_{m_{i}}^{\top }\nabla_{m_{i}}Q_{\phi}(x)\nabla_{m_{i}}Q_{\phi}(x)^{\top}v_{m_{i}}\right).$ (118) $$

We now take derivatives with respect to the remaining parameters. Writing the critic in terms of a sum over final layer weights:

$$ $\displaystyle Q_{\phi}(x)=\sum_{j=0}^{k-1}w_{j}\textrm{LayerNorm}_{j}[f_{M}(x)],$ (119) $$

we apply the chain rule to find the derivative of the critic with respect to each parameter vector $m_{i}\in M$:

$$ $\displaystyle\nabla_{m_{i}}Q_{\phi}(x)=\sum_{j=0}^{k-1}w_{j}\partial_{f_{i}} \textrm{LayerNorm}_{j}[f_{M}(x)]\sigma^{\prime}(m_{i}^{\top}x)x,$ (120) $$

where $\sigma^{\prime}$ denotes the derivative of the activation function which is bounded almost surely. Now, as $f_{M}(x)$ is Lipschitz and defined over a bounded set of parameters $\Phi$ and inputs $\mathcal{X}$ under 2, it follows that it must be a bounded function. We can thus apply the derivative bound to $\partial_{f_{i}}\textrm{LayerNorm}_{j}[f_{M}(x)]$ from Proposition 1, yielding:

$$ $\displaystyle\nabla_{m_{i}}Q_{\phi}(x)^{\top}v_{m_{i}}$ $\displaystyle=\sum_{j=0}^{k-1}w_{j}\mathcal{O}\left(\frac{\mathbbm{1}(i=j)- \frac{1}{k}}{\sqrt{k}}\right)x^{\top}v_{m_{i}},$ (121) $\displaystyle=\mathcal{O}\left(\frac{1}{\sqrt{k}}\right).$ (122) $$

Finally, we substitute into Eq. 118 to yield our desired result:

$$ $\displaystyle\frac{1}{\lVert v\rVert^{2}}\sum_{i=0}^{k-1}\left(\gamma v_{m_{i} }^{\top}\nabla_{m_{i}}Q_{\phi}(x^{\prime})\nabla_{m_{i}}Q_{\phi}(x)^{\top}v_{m _{i}}-v_{m_{i}}^{\top}\nabla_{m_{i}}Q_{\phi}(x)\nabla_{m_{i}}Q_{\phi}(x)^{\top }v_{m_{i}}\right),$ (123) $\displaystyle=\mathcal{O}\left(\frac{1}{\sqrt{k}}\right)\frac{1}{\lVert v \rVert^{2}}\sum_{i=0}^{k-1}\lVert v_{m_{i}}\rVert^{2}\left(\gamma\mathcal{O} \left(1\right)-\mathcal{O}\left(1\right)\right),$ (124) $\displaystyle=\mathcal{O}\left(\frac{1}{\sqrt{k}}\right)\frac{1}{\lVert v \rVert^{2}}\sum_{i=0}^{k-1}\lVert v_{m_{i}}\rVert^{2},$ (125) $\displaystyle=\mathcal{O}\left(\frac{1}{\sqrt{k}}\right)\frac{\lVert v_{m} \rVert^{2}}{\lVert v\rVert^{2}}=\mathcal{O}\left(\frac{1}{\sqrt{k}}\right),$ (126) $\displaystyle\implies$ $\displaystyle\gamma v^{\top}\nabla_{\phi}Q_{\phi}(x^{\prime})\nabla_{\phi}Q_{ \phi}(x)^{\top}v-(v^{\top}\nabla_{\phi}Q_{\phi}(x))^{2}\leq\left(\frac{\gamma \lVert v_{w}\rVert}{2}\right)^{2}+\mathcal{O}\left(\frac{1}{\sqrt{k}}\right).$ (127) $$

∎

###### Lemma 4 (Mitigating Nonlinear Instability) .

Under 2, using the LayerNorm $Q$-function defined in Eq. 8:

$$ $\displaystyle\frac{\left\lvert\left(r+\gamma Q_{\phi}(x^{\prime})-Q_{\phi}(x) \right)v^{\top}\nabla_{\phi}^{2}Q_{\phi}(x)v\right\rvert}{\lVert v\rVert^{2}}= \mathcal{O}\left(\frac{1}{\sqrt{k}}\right),$ (128) $$

almost surely for any test vector and any state-action transition pair $x,x^{\prime}\in\mathcal{X}$.

###### Proof.

We start by bounding the TD error, second order derivative and test vector separately:

$$ $\displaystyle\left\lvert\left(r+\gamma Q_{\phi}(x^{\prime})-Q_{\phi}(x)\right) v^{\top}\nabla_{\phi}^{2}Q_{\phi}(x)v\right\rvert$ $\displaystyle\leq\left\lvert\left(r+\gamma Q_{\phi}(x^{\prime})-Q_{\phi}(x) \right)\right\rvert\left\lVert\nabla_{\phi}^{2}Q_{\phi}(x)\right\rVert\left \lVert v\right\rVert^{2}.$ (129) $$

By the definition of the LayerNorm $Q$-function:

$$ $\displaystyle\lvert Q_{\phi}(x)\rvert$ $\displaystyle\leq\left\lVert w\right\rVert\left\lVert\textrm{LayerNorm}[f_{M}( x)]\right\rVert,$ (130) $\displaystyle\leq\left\lVert w\right\rVert\left\lVert\textrm{LayerNorm}[f_{M}( x)]\right\rVert,$ (131) $\displaystyle\leq\left\lVert w\right\rVert.$ (132) $$

where the final line follows from Proposition 1. As the reward is bounded by definition and $\left\lVert w\right\rVert$ is bounded under 2, we can bound the TD error vector as:

$$ $\displaystyle\left\lvert\left(r+\gamma Q_{\phi}(x^{\prime})-Q_{\phi}(x)\right) \right\rvert=\mathcal{O}(1),$ (133) $$

hence:

$$ $\displaystyle\frac{\left\lvert\left(r+\gamma Q_{\phi}(x^{\prime})-Q_{\phi}(x) \right)v^{\top}\nabla_{\phi}^{2}Q_{\phi}(x)v\right\rvert}{\lVert v\rVert^{2}}$ $\displaystyle\leq\left\lVert\nabla_{\phi}^{2}Q_{\phi}(x)\right\rVert.$ (134) $$

Our result follows immediately by using Lemma 5 to bound the second order derivative:

$$ $\displaystyle\frac{\left\lvert\left(r+\gamma Q_{\phi}(x^{\prime})-Q_{\phi}(x) \right)v^{\top}\nabla_{\phi}^{2}Q_{\phi}(x)v\right\rvert}{\lVert v\rVert^{2}}= \mathcal{O}\left(\frac{1}{\sqrt{k}}\right).$ (135) $$

∎

###### Lemma 5 .

Let 2 hold.
$\lVert\nabla_{\phi}^{2}Q_{\phi}(x)\rVert=\mathcal{O}\left(\frac{1}{\sqrt{k}}
\right).$

###### Proof.

We denote the partial derivative with respect to the $i,j$-th matrix elementas: $\partial_{m_{i,j}}\textnormal{{LayerNorm}}_{l}[f_{M}(x)]$. From Proposition 1 it follows:

$$ $\displaystyle\partial_{m_{i,j}}\textnormal{{LayerNorm}}_{l}\left[f_{M}(x)\right]$ $\displaystyle=\mathcal{O}\left(k^{-\frac{1}{2}}\left(\mathbbm{1}(l=i)+\frac{1} {k}\right)\right),$ (136) $\displaystyle\partial_{m_{s,t}}\partial_{m_{i,j}}\textnormal{{LayerNorm}}_{l}[ f_{M}(x)]$ $\displaystyle=\mathcal{O}\left(k^{-\frac{3}{2}}\left(\mathbbm{1}(l=i)+\mathbbm {1}(s=i)+\mathbbm{1}(l=s)+\frac{1}{k}\right)\right).$ (137) $$

We use these results find the partial derivatives of the LayerNorm $Q$-function:

$$ $\displaystyle\partial_{w_{u}}\partial_{m_{i,j}}Q_{\phi}(x)=\partial_{w_{u}} \sum_{l=0}^{k-1}w_{l}\partial_{m_{i,j}}\textrm{LayerNorm}_{l}[f_{M}(x)],$ (138) $\displaystyle=\partial_{w_{u}}\sum_{l=0}^{k-1}w_{l}\mathcal{O}\left(k^{-\frac{ 1}{2}}\right)\left(\mathbbm{1}(l=i)+\mathcal{O}\left(\frac{1}{k}\right)\right),$ (139) $\displaystyle=\mathcal{O}\left(k^{-\frac{1}{2}}\right)\left(\mathbbm{1}(u=i)+ \mathcal{O}\left(\frac{1}{k}\right)\right),$ (140) $\displaystyle\partial_{m_{s,t}}\partial_{m_{i,j}}Q_{\phi}(x)=\sum_{l=0}^{k-1}w _{l}\partial_{m_{s,t}}\partial_{m_{i,j}}\textrm{LayerNorm}_{l}[f_{M}(x)],$ (141) $\displaystyle=\sum_{l=0}^{k-1}w_{l}\mathcal{O}\left(k^{-\frac{3}{2}}\right) \left((\mathbbm{1}(l=i)+\mathbbm{1}(s=i)+\mathbbm{1}(l=s)+\mathcal{O}\left( \frac{1}{k}\right)\right),$ (142) $\displaystyle=\mathcal{O}\left(k^{-\frac{1}{2}}\right)\left(\mathbbm{1}(s=i)+ \mathcal{O}\left(\frac{1}{k}\right)\right).$ (143) $$

Now, from the definition of the Matrix 2-norm:

$$ $\displaystyle\lVert\nabla_{\phi}^{2}Q_{\phi}(x)\rVert=\sup_{v}\frac{v^{\top} \nabla_{\phi}^{2}Q_{\phi}(x)v}{v^{\top}v},$ (144) $$

for any test vector with $\lVert v\rVert_{\textrm{max}}<\infty$.
As the $Q$-function is linear in $w$, we can ignore second order derivatives with respect to elements of $w$ as their value is zero. The matrix norm can then be written in terms of the partial derivatives of $Q_{\phi}(x)$ as:

$$ $\displaystyle\lVert\nabla_{\phi}^{2}Q_{\phi}(x)\rVert=$ $\displaystyle\sup_{v}\frac{1}{v^{\top}v}\Bigg{(}\sum_{i=0}^{k-1}\sum_{j=0}^{d- 1}\sum_{u=0}^{k-1}v_{m_{i,j}}\partial_{w_{u}}\partial_{m_{i,j}}Q_{\phi}(x)v_{w _{u}}$ (145) $\displaystyle+\sum_{i=0}^{k-1}\sum_{j=0}^{d-1}\sum_{s=0}^{k-1}\sum_{t=0}^{d-1} v_{m_{i,j}}\partial_{m_{s,t}}\partial_{m_{i,j}}Q_{\phi}(x)v_{m_{s,t}}\Bigg{)}.$ (146) $$

Substituting for the bounds from Eq. 140 and Eq. 143 yields:

$$ $\displaystyle\lVert\nabla_{\phi}^{2}Q_{\phi}(x)\rVert$ $\displaystyle=\mathcal{O}\left(k^{-\frac{1}{2}}\right)\sup_{v}\frac{1}{v^{\top }v}\Bigg{(}\sum_{i=0}^{k-1}\sum_{j=0}^{d-1}\sum_{u=0}^{k-1}v_{m_{i,j}}v_{w_{u} }\left(\mathbbm{1}(u=i)+\mathcal{O}\left(\frac{1}{k}\right)\right)$ (147) $\displaystyle+\sum_{i=0}^{k-1}\sum_{j=0}^{d-1}\sum_{s=0}^{k-1}\sum_{t=0}^{d-1} v_{m_{i,j}}v_{m_{s,t}}\left(\mathbbm{1}(i=s)+\mathcal{O}\left(\frac{1}{k} \right)\right)\Bigg{)},$ (148) $$

Now, $v_{m_{i,j}}v_{w_{u}}=\mathcal{O}(\mathbbm{1}(\lvert v_{m_{i,j}}v_{w_{u}}\rvert
>0))$ and $v_{m_{i,j}}v_{m_{s,t}}=\mathcal{O}(\mathbbm{1}(\lvert v_{m_{i,j}}v_{m_{s,t}}
\rvert>0))$ hence:

$$ $\displaystyle\lVert\nabla_{\phi}^{2}Q_{\phi}(x)\rVert$ $\displaystyle=\mathcal{O}\left(k^{-\frac{1}{2}}\right)\sup_{v}\frac{1}{v^{\top }v}\Bigg{(}\sum_{i=0}^{k-1}\sum_{j=0}^{d-1}\sum_{u=0}^{k-1}\mathbbm{1}(\lvert v _{m_{i,j}}v_{w_{u}}\rvert>0)\left(\mathbbm{1}(u=i)+\mathcal{O}\left(\frac{1}{k }\right)\right)$ (149) $\displaystyle+\sum_{i=0}^{k-1}\sum_{j=0}^{d-1}\sum_{s=0}^{k-1}\sum_{t=0}^{d-1} \mathbbm{1}(\lvert v_{m_{i,j}}v_{m_{s,t}}\rvert>0)\left(\mathbbm{1}(i=s)+ \mathcal{O}\left(\frac{1}{k}\right)\right)\Bigg{)},$ (150) $\displaystyle=\mathcal{O}\left(k^{-\frac{1}{2}}\right)\sup_{v}\frac{1}{v^{\top }v}\Bigg{(}\sum_{u=0}^{k-1}\mathcal{O}\left(\mathbbm{1}(\lvert v_{w_{u}}\rvert >0)\right)+\sum_{i=0}^{k-1}\sum_{j=0}^{d-1}\mathcal{O}\left(\mathbbm{1}(\lvert v _{m_{i,j}}\rvert>0)\right)\Bigg{)},$ (151) $$

Finally, we notice that:

$$ $\displaystyle\frac{1}{v^{\top}v}=\mathcal{O}\left(\frac{1}{\sum_{u=0}^{k-1} \mathbbm{1}(\lvert v_{w_{u}}\rvert>0)+\sum_{i=0}^{k-1}\sum_{j=0}^{d-1}\mathbbm {1}(\lvert v_{m_{i,j}}\rvert>0)}\right),$ (152) $$

hence

$$ $\displaystyle\lVert\nabla_{\phi}^{2}Q_{\phi}(x)\rVert$ $\displaystyle=\mathcal{O}\left(k^{-\frac{1}{2}}\right)\sup_{v}\mathcal{O}\left (\frac{\sum_{u=0}^{k-1}\mathbbm{1}(\lvert v_{w_{u}}\rvert>0)+\sum_{i=0}^{k-1} \sum_{j=0}^{d-1}\mathbbm{1}(\lvert v_{m_{i,j}}\rvert>0)}{\sum_{u=0}^{k-1} \mathbbm{1}(\lvert v_{w_{u}}\rvert>0)+\sum_{i=0}^{k-1}\sum_{j=0}^{d-1}\mathbbm {1}(\lvert v_{m_{i,j}}\rvert>0)}\right),$ (153) $\displaystyle=\mathcal{O}\left(k^{-\frac{1}{2}}\right)\sup_{v}\mathcal{O}\left (1\right),$ (154) $\displaystyle=\mathcal{O}\left(\frac{1}{\sqrt{k}}\right).$ (155) $$

∎

###### Proposition 1 .

###### Proof.

###### Theorem 2 .

###### Proof.

###### Lemma 2 .

###### Proof.

###### Lemma 3 (Mitigating Off-policy Instability) .

###### Proof.

###### Lemma 4 (Mitigating Nonlinear Instability) .

###### Proof.

###### Lemma 5 .

###### Proof.

## Appendix C Implementation

Figure: Algorithm 2 Sample Efficient PQN with $Q(\lambda)$ (variables $s,a,s^{\prime},r,terminal,target(s)$ are batched)

## Appendix D Experimental Setup

All experimental results are shown as mean of 10 seeds, except in Atari Learning Environment (ALE) where we followed a common practice of reporting 3 seeds. They were performed on a single NVIDIA A40 by jit-compiling the entire pipeline with Jax in the GPU, except for the Atari experiments where the environments run on an AMD 7513 32-Core Processor. Hyperparameters for all experiments can be found in Appendix F. We used the algorithm proposed in present in appendix C. All experiments used Rectified Adam optimiser Liu et al. (). We didn’t find any improvements in scores by using RAdam instead of Adam, but we found it more robust in respect to the epsilon parameter, simplifying the tuning of the optimiser.

#### Baird’s Counterexample

For these experiments, we use the code provided as solutions to the problems of (Sutton & Barto, ) (^2^22[https://github.com/vojtamolda/reinforcement-learning-an-introduction/tree/main](https://github.com/vojtamolda/reinforcement-learning-an-introduction/tree/main)). We use a single-layer neural network with a hidden size of 16 neurons, with normalisation between the hidden layer and the output layer. To not include additional parameters and completely adhere to theory, we don’t learn affine tranformation parameters in these experiments, which rescale the normalised output by a factor $\gamma$ and add a bias $\beta$. However, in more complex experiments we do learn these parameters.

#### DeepSea

For these experiments, we utilised a simplified version of Bootstrapped-DQN (Osband et al., ), featuring an ensemble of 20 randomly initialised policies, each represented by a two-layered MLP with middle-layer normalisation. We did not employ target networks and updated all policies in parallel by sampling from a shared replay buffer. We adhered to the same parameters for Bootstrapped-DQN as presented in Osband et al. ().

#### MinAtar

We used the vectorised version of MinAtar (Young & Tian, ) present in Gymnax and tested PQN against PPO in the 4 available tasks: Asterix, SpaceInvaders, Freeway and Breakout. PQN and PPO use both a Convolutional Network with 16 filters with a 3-sized kernel (same as reported in the original MinAtar paper) followed by a 128-sized feed-forward layer. Results in MinAtar are reported in Fig. 8. Hyperparameters were tuned for both PQN and PPO.

#### Atari

We use the vectorised version of ALE provided by Envpool for a preliminary evaluation of our method. Given that our main baseline is the CleanRL (Huang et al., ) implementation of PPO (which also uses Envpool and Jax), we used its environment and neural network configuration. This configuration is also used in the results reported in the original Rainbow paper, allowing us to obtain additional baseline scores from there. Aitchison et al. (Aitchison et al., ) recently found that the scores obtained by algorithms in 5 of the Atari games have a high correlation with the scores obtained in the entire suite, and that 10 games can predict the final score with an error lower than 10%. This is due to the high level of correlation between many of the Atari games. The results we present for PQN are obtained by rolling out a greedy-policy in 8 separate parallel environments during training, which is more effective than stopping training to evaluate on entire episodes, since in Atari they can last hundreds of thousands of frames. We did not compare with distributed methods like Ape-X and R2D2 because they use an enormous time-budget (5 days of training per game) and frames (almost 40 Bilions), which are outside our computational budget. We comment that these methods typically ignore concerns of sample efficiency. For example Ape-X (Horgan et al., ) takes more than 100M frames to solve Pong, the easiest game of the ALE, which can be solved in few million steps by traditional methods and PQN.

#### Craftax

We follow the same implementation details indicated in the original Craftax paper Matthews et al. (). Our RNN implementation is the same as the MLP one, with an addtional LSTM layer before the last layer.

#### Hanabi

We used the Jax implementation of environments present in JaxMARL. Our model doesn’t use RNNs in this task. From all the elements present in the R2D2-VDN presented in Hu et al. (), we only used the duelling architecture Wang et al. (). Presented results of PQN are average across 100k test games.

#### Smax

We used the same RNN architecture of QMix present in JaxMARL, with the only difference that we don’t use a replay buffer, with added normalisation and $Q(\lambda)$. We evaluated with all the standard SMAX maps excluding the ones relative to more than 20 agents, because they could not be run with traditional QMix due to memory limitations.

#### Overcooked

We used the same CNN architecture of VDN present in JaxMARL, with the only difference that we don’t use a replay buffer, with added normalisation and $Q(\lambda)$.

## Appendix E Further Results

Figure: (a) Baird’s C.example
Refer to caption: x18.png

Figure: (a) CartPole:Training Time
Refer to caption: x22.png

Figure: Figure 8: Results in MinAtar
Refer to caption: x26.png

Figure: (a) Atari-Pong: PQN training
Refer to caption: x27.png

**Table 3: Scores in ALE.**
| Method (Frames) | Time | Gradient | Atari-10 | Atari-57 | Atari-57 | Atari-57 |
| --- | --- | --- | --- | --- | --- | --- |
|  | (hours) | Steps | Score | Median | Mean | >Human |
| PPO (200M) | 2.5 | 700k | 165 |  |  |  |
| PQN (200M) | 1 | 700k | 191 |  |  |  |
| PQN (400M) | 2 | 1.4M | 243 | 245 | 1440 | 40 |
| Rainbow (200M) | 100 | 50M | 239 | 230 | 1461 | 43 |

Figure: Figure 10: IQM computed over 3 seeds when training PQN with the ALE configuration proposed by Dopamine Castro et al. (). This configuration incorporates sticky actions and doesn’t set the done flag when an agent loses a life. With this setup, PQN can still outperform Rainbow, but it requires significantly more compute time (almost 5 hours), corresponding to 800 million frames, indicating a loss of sample efficiency. Sample efficiency might be recovered in this configuration by using a larger network or tuning the hyperparameters, but we leave this as future work. PQN is still much faster to train than a Dopamine agent, which requires multiple days depending on the hardware.
Refer to caption: x30.png

Figure: Figure 11: Improvement of PQN over Rainbow. Results refer to PQN trained for 400M frames, i.e. 2 hours of GPU time.
Refer to caption: x31.png

Figure: (a) Amidar
Refer to caption: extracted/5948673/Figures/experiments/atari-10/Amidar-v5.png

Figure: Figure 13: PQN learns a policy from an almost random policy. To further test PQN’s off-policiness, we conducted an experiment on Atari10 games using a highly random policy for collecting data, gradually shifting from 100% random to 70% random during training. As expected, the resulting policy was less effective compared to one with more exploitation. However, the key finding is that PQN can still learn a policy even when off-policiness is extremely high — i.e., when data is collected almost randomly from the environment, without following the learning policy.
Refer to caption: x32.png

Figure: (a) Craftax with MLP networks
Refer to caption: x33.png

Figure: Figure 15: Results in Smax
Refer to caption: x36.png

Figure: Figure 16: The buffer size scales quadratically in respect to the number of agents in SMAX.
Refer to caption: x37.png

Figure: Figure 17: Results in Overcooked
Refer to caption: x38.png

## Appendix F Hyperparameters

**Table 4: Craftax RNN Hyperparameters**
| Parameter | Value |
| --- | --- |
| NUM_ENVs | 1024 |
| NUM_STEPS | 128 |
| EPS_START | 1.0 |
| EPS_FINISH | 0.005 |
| EPS_DECAY | 0.1 |
| NUM_MINIBATCHES | 4 |
| NUM_EPOCHS | 4 |
| NORM_INPUT | True |
| NORM_TYPE | "batch_norm" |
| HIDDEN_SIZE | 512 |
| NUM_LAYERS | 1 |
| NUM_RNN_LAYERS | 1 |
| ADD_LAST_ACTION | True |
| LR | 0.0003 |
| MAX_GRAD_NORM | 0.5 |
| LR_LINEAR_DECAY | True |
| REW_SCALE | 1.0 |
| GAMMA | 0.99 |
| LAMBDA | 0.5 |

**Table 5: Atari Hyperparameters**
| Parameter | Value |
| --- | --- |
| NUM_ENVs | 128 |
| NUM_STEPS | 32 |
| EPS_START | 1.0 |
| EPS_FINISH | 0.001 |
| EPS_DECAY | 0.1 |
| NUM_EPOCHS | 2 |
| NUM_MINIBATCHES | 32 |
| NORM_INPUT | False |
| NORM_TYPE | "layer_norm" |
| LR | 0.00025 |
| MAX_GRAD_NORM | 10 |
| LR_LINEAR_DECAY | False |
| GAMMA | 0.99 |
| LAMBDA | 0.65 |

**Table 6: SMAX Hyperparameters**
| Parameter | Value |
| --- | --- |
| NUM_ENVs | 128 |
| MEMORY_WINDOW | 4 |
| NUM_STEPS | 128 |
| HIDDEN_SIZE | 512 |
| NUM_LAYERS | 2 |
| NORM_INPUT | True |
| NORM_TYPE | "batch_norm" |
| EPS_START | 1.0 |
| EPS_FINISH | 0.01 |
| EPS_DECAY | 0.1 |
| MAX_GRAD_NORM | 1 |
| NUM_MINIBATCHES | 16 |
| NUM_EPOCHS | 4 |
| LR | 0.00025 |
| LR_LINEAR_DECAY | True |
| GAMMA | 0.99 |
| LAMBDA | 0.85 |

**Table 7: Overcooked Hyperparameters**
| Parameter | Value |
| --- | --- |
| NUM_ENVs | 64 |
| NUM_STEPS | 16 |
| HIDDEN_SIZE | 512 |
| NUM_LAYERS | 2 |
| NORM_INPUT | False |
| NORM_TYPE | "layer_norm" |
| EPS_START | 1.0 |
| EPS_FINISH | 0.2 |
| EPS_DECAY | 0.2 |
| MAX_GRAD_NORM | 10 |
| NUM_MINIBATCHES | 16 |
| NUM_EPOCHS | 4 |
| LR | 0.000075 |
| LR_LINEAR_DECAY | True |
| GAMMA | 0.99 |
| LAMBDA | 0.5 |

**Table 8: Hanabi Hyperparameters**
| Parameter | Value |
| --- | --- |
| NUM_ENVS | 1024 |
| NUM_STEPS | 1 |
| TOTAL_TIMESTEPS | 1e10 |
| HIDDEN_SIZE | 512 |
| N_LAYERS | 3 |
| NORM_TYPE | layer_norm |
| DUELING | True |
| EPS_START | 0.01 |
| EPS_FINISH | 0.001 |
| EPS_DECAY | 0.1 |
| MAX_GRAD_NORM | 0.5 |
| NUM_MINIBATCHES | 1 |
| NUM_EPOCHS | 1 |
| LR | 0.0003 |
| LR_LINEAR_DECAY | False |
| GAMMA | 0.99 |

Table: Table 9: ALE Scores: Rainbow vs PQN (400M frames)

## Peer Reviews

confidence: 3
contribution: 3
initial_rating: 6
presentation: 3
questions: - Is PPO using the same number of parallel environments as PQN in all experiments? I couldn’t find this information in the paper. Could you share this information and add it to the paper’s revision?
- In line 428, the authors refer to a histogram in Appendix E, but there is no histogram. Do they mean the bar plot in Figure 12?
- Typically, layer norm is not applied to the post-activation but instead to the pre-activation. This is an important distinction. The theory still works with the preactivation layer norm as long as you don’t use activation functions that scale up the inputs.
soundness: 3
strengths: The work on simplifying deep reinforcement learning and removing techniques that might not be necessary, like replay buffer and target networks, is undoubtedly fundamental to deep reinforcement research. This has vast implications for rethinking the widely existing deep RL approaches and can help in other important directions, like scaling RL with the number of parameters/samples. This paper provides a unique view that challenges existing beliefs on the importance of replay buffers and target networks. The approach is effective and efficient since it can be implemented parallelized on GPU, which outperforms other baselines with respect to wall clock time.
summary: Modern deep-reinforcement learning resorts to techniques such as replay buffers and target networks to provide stability with nonlinear off-policy learning. However, learning becomes unstable without a replay buffer or target networks and can diverge. Recently, several works suggested using layer normalization or layer normalization in addition to l2 regularization to remedy this learning instability issue. This paper theoretically studies layer normalization’s role and identifies how layer normalization helps with stability and convergence. The paper also proposed a new method called PQN that uses layer normalization and parallelized environments to stabilize learning. The authors show the effectiveness of their method through a series of experiments on different domains of environments.
weaknesses: - The theoretical part of the manuscript is largely incoherent. 
  - The current manuscript scatters many things in the theory parts. It lacks a proper flow of ideas when describing the theoretical results and their implications, which makes it difficult to follow. Currently, it reads as bullet points, listing findings quickly without proper linking between subsequent findings or results.
  - For example, the current theorems and lemmas are not well integrated with the text before and after them. They read as detached components, making reading unnecessarily harder.
  - Notations can be improved. I suggest following Sutton & Barto (2018).
  - Two things that are instrumental for the results based on Jacobian analysis: inequality 3 and inequality 4. More discussion is needed to understand these two conditions and their implications.
  - Off-policy instability and nonlinear instability require their own theorem statements and separate proofs (even if previous works have shown them). In addition, the TD stability criterion needs a theorem statement about contraction mapping. 
  - The connection for why l2 is needed was not clear. It was directly introduced after layer norm without proper linking.
- PQN solves a problem different from the baseline DQN, but this was never discussed
  - The authors emphasize that PQN does not use a replay buffer or target network as an advantage over other methods, which is great. However, a similar emphasis is needed for the fact that PQN requires parallel environments and probably may fail if a single environment were used (the setting of the other baselines). Additionally, PQN solves another orthogonal problem (parallelized worlds) to the original RL problem (single world).
  - Figure 1 is inaccurate. The replay buffer is part of the agent, not an external component. This needs to be fixed.
  - The difference between Distributed DQN and PQN is unclear from Figure 1, although both solve the same problem (parallelized worlds), especially the point on synchronism and GPU is not clear.
- The paper has several inaccuracies.
  - The authors claim that their PQN is based on Peng’s Q($\lambda$), but the actual algorithm does not use eligibility traces. Instead, the authors use Q-learning with $\lambda$-return. This needs to be corrected. Additionally, the equation (line 326) needs to be derived from the first principle to make the paper accessible to the unfamiliar reader.
  - Algorithm 1 is Q-learning with one-step targets. The authors mention that they use $\lambda$-return target, so Algorithm 1 needs to be replaced with the algorithm actually used (Algorithm 2 in Appendix C).
  - The authors claim they stabilize learning Baird’s counterexample and use Figure 7a to demonstrate that. However, in Figure 7a, the error increases from the starting point until it plateaus. I don’t think meaningful learning has happened since the error has increased instead of decreased. I see that with layer norm or layer norm + L2 regularization, the error doesn’t increase without bound but at the same time, the problem is not solved either.
- Unfair or unclear empirical evaluation
  - The authors compare algorithms that work with parallel environments against ones that do not, which requires careful experiments to compare them. 
  - In Figure 3, the authors need to write rainbow (200M) or DQN (200M) to understand what those horizontal lines represent clearly.
  - When the authors say that PQN was trained for 400M frames, does that include the parallel environments? For example, if 128 parallel environments are used (according to Table 5), does this mean 3.125M frames were collected from each environment, resulting in a total of 400M frames, or does it mean that 400M frames were collected from each environment, resulting in 128x400M=51200M overall frames? The first option gives a fair comparison, but the second option is biased towards PQN since significantly more experience is used. I would like the authors to clarify this point.
  - The authors used only 3 independent runs for Atari, relying on precedence.  This is a too low number to provide any statistical significance. Even if something was accepted before, it does not mean it is correct. I highly suggest the authors increase the number of independent runs to at least 10. This should be possible since both PQN and PPO are efficient (small clock time) and easy to run, according to the paper's claims.
  - The authors mentioned that DQN/Rainbow uses 50M updates compared to 700k updates for PQN. I think it is unclear how these numbers are obtained, especially for PQN.
  - In Figure 6a, PQN still learns well without divergence when no layer normalization is used. What is the reason? Why has no divergence happened?
  - Since the authors compare against DQN and Rainbow (methods that use a single environment), an ablation where a different number of environments are considered (e.g., n=1 and n=10) to understand the provided stability coming from parallelized environments. I think parallel environments make the gradient signal more reliable compared to the single environment case, which is more prone to noisy gradients. This may be instrumental for PQN; thus, an ablation is needed.

&nbsp;

**Minor issues:**

- line 9 in Algorithm 1 and line 10 in Algorithm 2 are incorrect. The negative sign should be a plus. Additionally, I think the use of $\texttt{StopGrad}$ can be eliminated with the TD error vector you gave in Eq. 2

&nbsp;
&nbsp;

Overall, I believe the ideas from this paper could serve as a good contribution to the community, but the current manuscript is not ready for publication and needs significant improvement. I'm willing to increase my score, given that the authors improved their manuscript's quality by 1) rewriting the theoretical part to make it coherent and more rigorous, 2) Fixing the inaccuracies, and 3) improving the empirical evaluation quality.

---

confidence: 3
contribution: 3
initial_rating: 8
presentation: 3
questions: - Including Baird's counterexample results in the main text would strengthen the paper in my opinion, by providing a clearer connection between the theoretical analysis and experimental validation.
- (Minor) The PPO baseline comparison in Figure 3 could be more consistent, though I understand the thinking to save compute. The paper could either compare both methods using 4e8 training frames in Figure 3(a), or include PPO results directly in Figure 3(b) across all Atari environments. Either approach would more effectively demonstrate PQN's competitiveness against established policy gradient methods, in terms of sample efficiency.
soundness: 4
strengths: - This paper is well-written and easy to follow, with clear presentation of both theoretical derivations and experimental results.
- The paper's motivation is clear and compelling enough to me: it mainly provides a simplified Q-learning baseline that effectively leverages GPU parallelization and vectorized environments.
- The proposed experimental evaluation is relatively comprehensive: it covers multiple domains including proof-of-concept environments, standard single-agent benchmarks such as Atari and Crafter, and multi-agent scenarios. They also covers variants of the Q learning methods to support the claim better in general.
summary: This paper proposes simplifications to multiple components of the Deep Q-Network (DQN)/TD learning method to enable more efficient training on a single GPU, offering a potential DQN baseline for future research. The modifications include:
- Eliminating target network tricks: The authors theoretically demonstrate that combining layer normalization with L2 regularization leads to convergent temporal difference (TD) learning. They then conducted experiments to validate this empirically by removing the target network update tricks.
- Removing replay buffer for experience replay: The paper identifies the replay buffer as a memory bottleneck that limits single-GPU training. While directly removing the replay buffer impacts sample efficiency, the paper demonstrates that when combined with vectorized environments, the GPU-based training method achieves better wall-clock time efficiency.
- Batch-wise rollout using vectorized environments: The paper implements DQN training in a batch-wise manner, leveraging GPU parallelization (after removing the replay buffer). Instead of single actor rollout, the paper uses vectorized environments to generate rollouts in batch.

To validate their approach, the authors conduct comprehensive experiments on both theoretical/proof-of-concept environments (Baird's counterexample) and standard benchmarks (Atari and Crafter). Their results show that the proposed Parallelized Q-Network (PQN) achieves comparable performance to well-known baselines like PPO and Rainbow. Through ablation studies, they further demonstrate the importance of network normalization and justify the removal of the replay buffer.
weaknesses: There is no major weakness of this paper, but feel free to check the question section for minor questions.

---

confidence: 2
contribution: 4
initial_rating: 8
presentation: 4
questions: One question I have is why does the Jacobian have to be negative definite on the unit Circle? Why is simple negative definiteness not enough? Since any vector can be written as || u || = c || v || where c is a positive constant, it seems that simply negative definiteness of the Jacobian is required? It would be helpful to the reader if the authors could give more justification and/or insight into the reason negative definiteness on unit circle is required, or the authors may wish to reconsider their results and see whether the restriction can be removed.
soundness: 3
strengths: The paper is well written, well formatted and quite readable. The authors present the essence of their results very well and show that stability of TD algorithms reduces to checking that a Jacobian is negative definite on the unit circle. This is a nice succinct and somehow intuitive result. Given the complexity of the proof, it was good to see the summary presented so concisely. 

Other results are then presented after this, including some insight into the causes of instability and then the approach using the layernorm to obtain stabilisation. The presentation here was not as clear as the above, but still acceptable and still concise. 

The authors claim that their parallelised version of Q learning is motivated well and this seems to be backed up by experiments.

The overall implications of the authors' results are very significant: they have discovered and captured the root cause for TD instability, they have proposed an improvement which guarantees instability and they have showed that their new PQN performs exceptionally well on some examples.
summary: This paper analyses stability in temporal difference (TD) methods. The main contributions are theoretical proofs that (i) TD instability can be established using a Jacobian evaluated on the unit circle; and (ii) using the Layernorm regularisation technique can ensure stability. This then leads the authors to propose a deep Q-learning algorithm called PQN which is comptetive with the PPO approach to reinforcement learning.
weaknesses: A criticism is that the proof of the main theoretical results is long (there are 20 pages of additional material) and I would say not particularly well organised. Before the authors give the proofs, in my view, it would be good for them to outline the main steps. I found the proofs hard to follow and as one goes through the proofs, there is a feeling of being somewhat adrift. In other words, the summary in the main paper is good; the actual proofs in the appendix are less well clear.

---

confidence: 3
contribution: 4
initial_rating: 8
presentation: 3
questions: See above
soundness: 3
strengths: 1. PQN is straightforward and easy to implement. It removes the need for a target network and simplifies TD Learning.
2. The paper includes theoretical analysis showing that regularization can help keep TD Learning stable without using target networks.
3. PQN achieves higher computational efficiency compared with baseline methods, with minimal impact in sample efficiency.
summary: This paper introduces Parallelised Q-Network (PQN), a streamlined deep online Q-Learning algorithm. PQN comprises two key components: a TD Learning objective without a target network, which instead applies layer normalization and L2 regularization, and a parallelized sampling approach that avoids the use of a replay buffer by leveraging vectorized environments. The authors provide theoretical analysis to support the claim that regularization can stabilize TD Learning in the absence of target networks. PQN demonstrates competitive performance across a wide range of environments, achieving results in significantly less wall-clock time.
weaknesses: 1. It’s somewhat counterintuitive that PQN maintains sample efficiency while training only on online samples without a replay buffer. Additional explanation would help readers understand this aspect better.
2. The removal of the target network in TD Learning and the parallelized sampling are independent components of the algorithm, yet their individual contributions to overall performance are unclear. More controlled experiments, like the one in Figure 6.d, would clarify the impact of each component.
3. The parallelized sampling approach depends on vectorized environments, which feels more like an engineering choice than a novel contribution and is not feasible in many real-world applications. A significant portion of the wall-clock savings seems to come from this aspect. If my understanding is correct, DQN could also potentially eliminate the replay buffer and use a similar parallelized sampling approach. It would be informative to see how DQN performs under this setup.

Meta-review: This paper makes two contributions. First, a proof that TD learning converges when the network uses layer normalization and weight-decay. This is demonstrated in experiments that show that one need not use a target value network or a “replay buffer” (there was a lot of back and forth on what constitutes a replay buffer with the reviewers). Second, the authors use vectorized environments for batched-rollouts to speed up training. We would like to thank the reviewers and the authors for a healthy discussion that has led to improvements in the manuscript during the review process.

I recommend that this paper be accepted. I encourage the authors to tie up some of the loose ends (e.g., public comments on this forum, as well as comments by Reviewer LgPr) in the camera-ready manuscript.

Decision: Accept (Spotlight)