Understanding, Predicting and Better Resolving Q-Value Divergence in Offline-RL
Authors: Yang Yue
Rui Lu
Bingyi Kang
Shiji Song
Gao Huang
Venue: NeurIPS 2023 poster (2023)
Keywords: Offline RL
Theory
TL;DR: Theoretically understand Q-value divergence in offline RL. A simple but effective solution is provided with machanism explained.

Abstract: The divergence of the Q-value estimation has been a prominent issue offline reinforcement learning (offline RL), where the agent has no access to real dynamics. Traditional beliefs attribute this instability to querying out-of-distribution actions when bootstrapping value targets. Though this issue can be alleviated with policy constraints or conservative Q estimation, a theoretical understanding of the underlying mechanism causing the divergence has been absent. In this work, we aim to thoroughly comprehend this mechanism and attain an improved solution. We first identify a fundamental pattern, \emph{self-excitation}, as the primary cause of Q-value estimation divergence in offline RL. Then, we propose a novel Self-Excite Eigenvalue Measure (SEEM) metric based on Neural Tangent Kernel (NTK) to measure the evolving property of Q-network at training, which provides an intriguing explanation of the emergence of divergence. For the first time, our theory can reliably decide whether the training will diverge at an early stage, and even predict the order of the growth for the estimated Q-value, the model's norm, and the crashing step when an SGD optimizer is used. The experiments demonstrate perfect alignment with this theoretic analysis. Building on our insights, we propose to resolve divergence from a novel perspective, namely improving the model's architecture for better extrapolating behavior. Through extensive empirical studies, we identify LayerNorm as a good solution to effectively avoid divergence without introducing detrimental bias, leading to superior performance. Experimental results prove that it can still work in some most challenging settings, i.e. using only 1$\%$ transitions of the dataset, where all previous methods fail. Moreover, it can be easily plugged into modern offline RL methods and achieve SOTA results on many challenging tasks. We also give unique insights into its effectiveness.

Understanding, Predicting and Better Resolving Q- Valuue Divergence in Offline-RL
Yang Yue 1 Rui Lu
Bingyi Kang* Shiji Song'
Gao Huang ByteDance Inc.
Department of Automation, BNRist, Tsinghua University
{1e-y22, r-lu21] (llu2)@mmalss @mails Ammalsstsinghun edu. en bingykang@gmail.
com com
{shijis, gaohuang @achuang)tssinguua @tsinghua. edu.
Abstract
The divergence of the Q-value estimation has been a prominent issue in offline reinforcement learning (offline RL), where the agent has no access to real dynamics. Traditional beliefs attribute this instability to querying out-of-distribution actions when bootstrapping value targets. Though this issue can be alleviated with pol- icy constraints or conservative Q estimation, a theoretical understanding of the underlying mechanism causing divergence has been absent. In this work, we aim to thoroughly comprehend this mechanism and attain an improved solution. We first identify a fundamental pattern, self-excitation, as the primary cause of Q-value estimation divergence in offline RL. Then, we propose novel Self-Excins Eigenvalue Measure (SEEM) metric based on Neural Tangent Kernel (NTK) measure the evolving property of Q-network training, which provides intrigu- ing explanation on the emergence of divergence. For the first time, our theory can reliably decide whether the training will diverge at early stage, and even predict the order of the growth for the estimated Q-value, the model's 's norm, and the crashing step when SGD optimizer is used. The experiments demonstrate perfect alignment with this theoretic analysis. Building on our insights, we propose resolve divergence from novel perspective, namely regularizing the neural net- work's generalization behavior. Through extensive empirical studies, we identify LayerNorm as good solution to effectively avoid divergence without introducing detrimental bias, leading to superior performance. Experimental results prove that it can still work in some most challenging settings, i.c. using only 1% transitions of the dataset, where all previous methods fail. Moreover, can be easily plugged into modern offline RL methods and achieve SOTA results on many challenging tasks. We also give unique insights into its effectiveness. Code can be found at www.ggittubioo.
Introduction
Off-policy Reinforcement Learning (RL) algorithms are particularly compelling for robot control [39; 13: 17] due to their high sample efficiency and strong performance. However, these methods often suffer divergence of value estimation, especially when value over-estimation becomes unmanageable. Though various solutions (e.g., double q- nentwwork [19]) proposed alleviate this issue in the online RL setting, this issue becomes more pronounced in offline RL, where the agent can only learn from offline datasets with online interactions prohibited [32]. As result, directly employing off-policy algorithms in offline settings confronts substantial issues related to value divergence [14; 30].
This raises natural yet crucial question: why is value estimation in offline RL prone to divergence, and how can we effectively address this issue? Conventionally, deadly triad [46; 4; 47; 48] is
"qual contribution. 'Coresspondig author.
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

identified to be the main cause of value divergence. Specifically, it points out that a RL algorithm susceptible to divergence at training when three components are combined, including off-policy learning, function approximation, and bootstrapping. As special case of off-policy RI., offline RL. further attributes the divergence to distributional shift: when training a Q-value function using the Bellman operator, the hootstrapping operation frequently queries the value of actions that are unseen in the dataset, resulting in cumulative extrapolation errors [14; 30; 33J. To mitigate this issue, existing algorithms either incorporate policy constraints between the learned policy and the behavior policy [42; 14; 53; 23; 30; 12; 51; 8], or make conservative/ensemble Q-value estimates [31; 38; 2: 41; 57; 15]. While these methods demonstrate some effectiveness by controlling the off- aff-policy degree and bootstrapping two components of the deadly triad -riy often overlook the aspect of function approximation. This neglect leaves several questions unanswered and certain limitations unaddressed in current methodologies.
How does divergence actually occur? Most existing studies only provide analyses of linear functions as Q-values, which is usually confined to contrived toy examples [47; 4]. The understanding and explanation for the divergence of non-linear neural networks in practice are still lacking. Little is known about what transpires in the model when its estimation inflates to infinity and what essential mechanisms behind the deadly triad contribute to this phenomenon. We are going to answer this question from the perspective of function approximation in the context of offline RL.
How to avoid detrimental bias? Despite the effectiveness of policy constraint methods in offline RL, they pose potential problems by introducing detrimental bias. Firstly, they explicitly force the policy to be near the behavior policy, which can potentially hurt performance. Also, the trade-off between performance and constraint needs to be balanced manually for each task. Secondly, these methods become incapable of dealing with divergence when more challenging scenario is encountered, e.g., the transitions are very scarce. Instead, we are interested in solutions without detrimental bias.
Our study offers novel perspective on the challenge of divergence in Offline-RL. We dissect the divergence phenomenon and identify seff-excitation-a process triggered by the gradient update of the Q-value estimation model as the primary cause. This occurs when the model's ' learning inadvertently elevates the target Q-value due its inherent generalization ability, which in turn encourages an even higher prediction value. This mirage-like property initiates a self-excitation cycle, leading to divergence. Based on this observation, we develop theoretical tools with Neural Tangent Kernel (NTK) to provide in-depth understanding and precise prediction of Q-value divergence. This the first work, to our knowledge, that accurately characterizes neural network dynamies in offline RL and explains the divergence phenomenon. More specifically, our contributions are three-fold:
We offer a detailed explanation of Q-value estimation divergence in offline RL settings with neural networks as non-linear estimators. We propose novel metric, the Self- Excite Eigenvalue Measure (SEEM), which is defined as the largest eigenvalue in the linear iteration during training, serving as an indicator of divergence. It also elucidates the fundamental mechanisms driving this divergence.
Explanation:
Prediction: We can foresee the divergence through SEEM an early stage of training before its estimation explodes. If divergence occurs, our theoretical framework is capable of predicting the growth pattern of the model's norm and the estimated Q-value. When the SGD optimizer used, we can even correctly predict the specific timestep at which the model is likely to crash. Our experiments show our theoretical analysis aligns perfectly with reality.
Effective Solution: Drawing from our findings, suggest mitigating divergence by regularizing the model's 's generalization behavior. This approach allows us to avoid imposing strict constraints on the learned policy, while still ensuring convergence. Specifically, viewed through the lens of SEEM, we find that MLP networks with LayerNorm exhibit excellent local generalization properties while MLPs without LayerNorm don't 't Then we conduct experiments demonstrate that with LayerNorm in critic network, modern offline RL algorithms can consistently handle challenging settings when the offline dataset is significantly small or suboptimal.
2 Preliminaries
Notation. A Markov Decision Process (MDP) is defined by tuple M (S, A, P,r,,, where
S is the state space and A C dRdd is the action space. P: S x A + A(S) is the transition dynamics mapping from state-action pair to distribution state space. r S A ++ R is the reward function and v E A(S) is the initial state distribution. An offline RL algorithm uses offline dataset

---

walker2d medium-22 seeds)
antmace medium
seeds)
NTK similarity
Action similarity
NTK similarity
Action simiiartty
1,0
0.8
0.8
64
0.4
62
0.2
0.2
60
200
500
750
1000
0.0
1000
1000 0.0
500
750
1000
0.0
253
500
750
1000
K Kaseps
Katepa
K steps
K Kalaga
Figure 1: NTK similarity and action similarity along the training in two offline D4RL tasks. Different colors represent different seeds. The cosine similarity between the current step and the last step is computed. We can see all runnings reach a steady NTK direction and policy actions after a time to.
D == {(s, ar, si,, r:)] which consists of M finite samples from interactions with MDP M. We write E D if there exists some E [M] such that 8;2 = 8, similar for notation (s,a) E D.
In this paper, we consider using I--layer ReL.U-activated MLP as the Q-value approximator: Each layer has parameter weight We R and bias by E Te The activation function is ReLU: (xx) = max(æ, 0). Integer de represents the dimensionality of the Een hidden layer, where do ds dA and dz =
since a scalar estimation is required for Q-value. Denote be the vector of all learnable parameters. We use Qo(s.a)
Vee W o. bu..... ,.. 1L,,y] by] fo a) to denote the estimated Q-value using neural network fo with parameter O. Also, an(s) = arg max, Qo a) is its induced policy. Sometimes we will write (X) E RM where X [(s1, a1),... (SM: ams)] stands for the concatenation of all inputs (s,, 02) E D in dataset.
Considering Q-learning glearning-based offline RL algorithm that alternates between the following two steps:
r-r(s, a) + max Qo(s', d'), 0 0 v/ (Q-(Qoss)))
(1)
where 77 is the learning rate for updating the Q network. Practical algorithms might use its actor-critic form for implementation, we stick to this original form for theoretical analysis. In the following analysis, call these two steps Q-value iteration.
Neural Tangent Kernel (NTK). NTK [21] is an important tool to analyze the dynamics and generalization of neural networks. Intuitively, it can be understood a special class of kernel function, which measures the similarity between & and x', with dependency on parameter 0. Given a neural network fo and two points æ, x' in the input space, the NTK is defined as kNTK (az, x') == (Vefe ()., fe (x') (z")), where the feature extraction is óg (x) == (e)). If we have two batches of input X, X', we can compute the Gram matrix as Go X, X') do XX) (oo(X))
Intuitively speaking. NTK measures the similarity between two inputs through the lens of a non-linear network fe. Such similarity can be simply understood as the correlation between the network's value predictions of these two inputs. If kNTK (x. x²) is a large value, it means these two points are "close" from the perspective of the current model's parameter. As a result, if the model's parameter changes near 0, say. from 0 to O' = 0 AO, denote the prediction change for x as A (x) = Se (x)- - ʃo (æ) and for z' as Af(=²) = for (z')= fo(x'). Ax( and Af(x') will be highly correlated, and the strength of the correlation is proportional to KNTK (x, x'). This intuition can be explained from simple Taylor expansion, which gives Af( = AO <e(x) and Af(x') AOT "e(x'). These two quantities will be highly correlated if de (x) and (a(")) are well-aligned, which corresponds to large NTK value between and æ'. To summary, kNTK (æ, x²) characterizes the subtle connection between fe (œ) and So (x') caused by neural network's 's generalization.
3 Theoretical Analysis
In this section, we will investigate the divergence phenomenon of Q-value estimation and theoretically explains it for a more comprehensive understanding. Note that in this section our analysis and experiments are conducted in a setting that does not incorporate policy constraints and exponential moving average targets. Although the setting we analyze has discrepancies with real practice, our analysis still provides valueable insight, since the underlying mechanism for divergence is the same. We first focus on how temporal difference (TD) error changes after single step of Q-value iteration Equation (1). Our result the following theorem.
3

---

Theorem 1. Suppose that the network's parameter at iteration t is 0,. For each transition
(8i, di, Tt) in dataset, denote r = - T'M]" E RRM, to, (s) = arg max. Qo, (8, a). Denote Ti, (8141, Ne, (8t+1)). Concatenate all , to be X; Denote un = fo. (X) (r + fo, (X;)) to be TD error vector at iteration t. The learning rate 1j is infinitesimal. When the maximal point of
Qo, is stable as I. increases, we have the evolving equation for Ut++1 as U (I + n)),)um, Ut, where
A, (x;)
- -00 (X)) 6o. (X) = -Go, (X; X) Go. (X, (X
G is defined in Section 2. Detailed proof is left in Appendix B. Theorem states that when the policy action fto, (s) that maximizes Q-values keep stable, leading to invariance of X", each Q-value iteration essentially updates the TD error vector by a matrix A,, which is determined by a discount factor y and the NTK between X and X
3.1 Linear Iteration Dynamics and SEEM
Although Theorem finds that TD update is a linear iteration under certain conditions, the dynamic is still complex in general. The policy ate (s) that maximizes Q-values may constantly fluctuate with O: over the course of training, leading to variations of X; Also, the kernel Go(-:) has dependency on parameter O. This causes non-linear dynamics of u, and the model parameter vector O. However, according to our empirical observation, we discover the following interesting phenomenon.
and Assumption 2. (Existence of Critical Point) There exists time step to. terminate kernel k(-,)) and stabilized policy state-action X* such that
(=,,0') |>||||||
k(æ.=²)
= o(1) and | i+1 x;
||
[(c] ||xl
o(n) for any x, x' ES x A when > fo- We also assume that | x [[.).
This assumption states that the training dynamics will reach a steady state, where the policy ite, stabilizes and NTK converges to ter- minate direction. Usually, constant NTK requires the width of the network to be infinite(21]. But in our experiments! the convergence of NTK is still observed in all our environments despite finite width (see Figure 1). Results of more environments appendix. It should be pointed out that NTK direction stability itself does not imply model convergence; rather, it only indicates that the model evolves in steady dynamics. Apart from NTK, we also observe that ne, converges after a certain timestep to. Moreover, we find that te. always shifts to the border of the action set with max norm (see Figure 15). The reason for the existence of this critical point is elusive and still unknown. We hypothesize that it originates from the structure of gradient flow ensuring every trajectory is near an attractor in the parameter space and ends up in steady state. Based on the critical point assumption, we prove that the further evolving dynamic after the critical point is precisely characterized by linear iteration as the following theorem. We also explain why such a stable state can persist in appendix.
1/Q step
attached in the
0.4
0.2
0.00144 x + 0.87
R-squared 1.00000
200
400
600
Steps
Figure 2: Prediction Ability Linear decay of inverse Q- value with SGD and L = 2.
use SGD optimizer for Q-value iteration with learning rate to be infinitesi- mal. Given iteration t >> to, and A A= ==1CC 'GiX" .X) - G X), where G is the Gram matrix under terminal kernel k. The divergence of uh is equivalent to whether there exists an eigenvalue of A becomes parallel
Theorem 3. Suppose
such that Re(A) > 0. If converge, we have = -I (nnAA( Aja-ta. Utg Otherwise, to the eigenvector the largest eigenvalue X of A, and its norm diverges to infinity at following order
(1
|| 1121122 =0 (0
(2)
(1-C'"
for some constant C' to be determined and L. is the number of layers of MLP. Specially, when T. = 2, it reduces to O (T=)m)
Theorem 3 predicts the divergence by whether Amax (A) is greater to O. We term this divergence detector Amax ()) as Self-Excite Eigenvalue Measure, or SEEM. Essentially, Theorem 3 states that
'To affirm the practicability of our analysis, we conduct experiments on widely used offline RL benchmark D4RL [10], and utilize 3 L-yee MLP instead of a toy example. More details can be found in the appendix.

---

walker2d--ediium- exxeeeeriiimeeee-2222222
antmaze-mediim-pleyy-e seed
arimaaze-medduu-ppaay-000 seed 2
0.8
0.3
divarge convergo
10
10
0.6-
10
Q Oovalu
0.2
0.4
0.4
10
10
0.1
0.2
10
0.2
10
0.0
0.0
0.0
990 0:22
033 190 1:26 03Z7 0225 6525 392 04
500
1000
0
500
1000
Figure 3: The divergence indication property of SEEM. The left figure shows the SEEM value with respect different discount factors 7, Theorem 3 states that larger y has larger SEEM value. The star point means the model's prediction eventually exceeds a large threshold (10%) in training, which means diverging. We can see that positive SEEM is perfectly indicative of the divergence. This is also true within the training process. From the middle and right figures, we can see that the prediction Q-value (in blue) is stable until the normalized kernel matrix's SEEM (in red) rises up to large positive value, then we can observe the divergence of the model.
we can monitor SEEM value to know whether the training will diverge. To validate this, we conduct experiments with different discount factor y. The first term in A is linear to y, thus larger y is more susceptible to diverging. This exactly what we find reality. In Figure 3, we can see that experiments with large y have positive SEEM, and the training eventually diverges (*). Moreover, SEEM can also faithfully detect the trend of divergence during the training course. On the right of Figure 3, we can see that the surge of SEEM value is always in sync with the inflation of estimation value. All these results corroborate our theoretical findings.
3.2 Prediction Ability
In Section 3.1. we have demonstrated that SEEM can reliably predict whether the training will diverge. In fact, our theory is able to explain and predict many more phenomena, which are all exactly observed empirically. Here we list some of them, and the detailed proofs are left in supplementary.
Terminating (Collapsing) Timestep. Theorem 3 predicts that with an SGD optimizer and 2-layer MLP, the inverse of Q-value decreases linearly along the timestep, implying terminating timestep T CrA The Q-value estimation will approach infinite very quickly beside T, and ends in singularity because the denominator becomes zero. From Figure 2, we can see that is true. The terminal Q-value prediction's 's inverse does decay linearly, and we can predict when it becomes zero to hit the singularity. Specifically, we use the data from 450th to 500th step to fit linear regression, obtaining the green dotted line 1/ 1/l|ue|l R 11.44 x 10 + 0.87, which predicts a singular point at 605 iteration. We then continue training and find that the model indeed collapses at the very point, whose predicted value and parameter become NaN.
Linear Norm Growth for Adam. While Theorem 3 studies SGD as the optimizer, Adam is more adopted in real practice. Therefore, we also deduce similar result for the Adam optimizer.
Theorem 4. Suppose we use Adam optimizer for Q-value iteration and all other settings are the same as Theorem 3. After >> to, the model will diverge if and only if Xmax >> O. If diverges, we have ||O,|| = nvpPt + o(t) and [|u, = ((t²) where P and L are the number of parameters and the number of layers for network fe. respectively.
Again, the detailed proof is left in the supplementary. Theorem 4 indicates that with Adam optimizer, the norm of the network parameters grows linearly and the predicted Q-value grows as a polynomial of degree along the time after a critical point fg- We verify this theorem in D4RL environments in Figure 4. We can see that the growth of the norm || exhibits straight line after a critic point tg. Moreover, we can see log Q 3ogg t + c, which means Q-value prediction grows cubically with time. Number 3 appears here because we use 3-layer MLP for value approximation. All these findings corroborate our theory, demonstrating our ability to accurately predict the dynamic of divergence in offline RI.
3.3 Discussions
Interpretation of SEEM and explain the self-excitation The intuitive interpretation for Amay (anxx(A) (A) is as below: Think about only one sample (so, ao) and its stabilized next step state-action
5

---

ne makor25-maddum22 (5 seeds)
artmaoedttddkmm wttp://wtt.ii..ppp----... (3 seeds)
Pgg(C) va. legit)
The
1gg legi))
80300
R-squared 0. 0.99778
60000
12
10300
12
40000
60300
20000
10
2.88
4433
10
Resquareed = 0.99619
5.00
5.25
6.00
K 750 1000
5.00
6.00
log(stops)
log(steps)
Figure 4: Prediction ability- Linear growth of network parameter's 's norm and polynomial growth of predicted Q-value with an Adam optimizer. Please note that in the 2th and 4th figures, both Q-value and steps have been taken the logarithm base 10. Different color represents different seeds.
xx (s1, a1). Without loss of generality, we assume fe, (a)) <r+tfe..,""") If SEEM is positive, we have Goo (a",, ,x"))) is larger than (æ, æ).. Recall that Go, (ax" depicts the strength of the bond between x and x" because of generalization. So we know that, whren updating the value of
fz(z) towards r + Tfo.. (x*), the Q-value iteration inadvertently makes fe.., (x*) increase even more than the increment of fe... (x). Consequently, the TD error r + +fe(== - fo( expands instead of reducing, due to the target value moving away faster than predicted value, which encourages the above procedure to repeat. This forms positive feedback loop and causes self-excitation. Such mirage-like property causes the model's parameter and its prediction value diverge.
Stability of policy action ãe, (s). Previous Theorem I and Theorem 3 rely on the stability of the policy action ne. (s) after to, and here we further elucidate why such stability occurs. Critic networks without normalization have tendency to output large values for extreme points, i.e., action Aez at the boundary of action space (see detailed explanation Section 4), making these extreme points easily become policy actions during optimization. We have observed that policy actions tend to gravitate toward extreme points when Q-value divergence takes place (see Figure 15), accounting for the stability of policy action tee (s). Consequently, as aforementioned, a loop is formed: Q(s,, a) keeps chasing Q(s', ,Ae)), ,, leading the model's parameter to diverge along certain direction to infinity.
Relations to Linear Setting. Since the linear regression model can be regarded as special case for neural networks, our analysis also directly applicable to linear settings by plugging in G(X, = X X, which reduce to study of deadly triad in linear settings. Therefore, our analysis fully contains linear case study and extends it to non-linear value approximation. In Appendix F.5, we present that our solution (Section 4) can solve the divergence Baird's Counterexample, which was first to show the divergence of Q-learning with linear approximation.
Similarities and Differences to Previous works. Our work shares similar observations of the con- nection between Q-value divergence and feature rank collapse with DR3 [29]. However, our work is different in the following aspects. First, DR3 attributes feature rank collapse to implicit regularization where L(0)) O, and it require the label noise from SGD. Actually, such near-zero critic loss assumption is mostly not true in the real practice of RL. Conversely, SEEM provides a mechanism behind feature rank collapse and Q-value divergence from the perspective of normalization- free network's 's pathological extrapolation behavior. It is applicable to more general settings and provides new information about this problem. Moreover, we formally prove in Theorems 3 and 4 that the model's parameter 0 evolves linearly for the Adam optimizer, and collapses at certain iter for SGD by solving an ODE, This accurate prediction demonstrates our precise understanding of neural network dynamics and is absent in previous work like DR3, Last, the solution to value divergence by DR3 involves searching hyperparameter tuning for co and also introduces extra computation when computing the gradient of á(s', a') to get Remp(() Our method (elaborated in Section 4) is free of these shortcomings by simple LayerNorm.
4 Reducing SEEM By Normalization
In Section 3, we have identified SEEM as a measure of divergence, with self-excitation being the primary catalyst for such divergence. In essence, a large SEEM value arises from the improper link between the dataset inputs and out-of-distiibution data points. To gain an intuitive grasp, we visualize the NTK value in simple 2-dimensional input space for a two-layer MLP in Figure 5. We designate a reference point mo = (..1, and calculate the NTK value Go( ar) for a range of x in the input domain. The heatmap displays high values at the boundaries of the input range.

---

nommalized NTK healmap wio Lh
10
normaized NIK heatmap with IN
100
05
2
2
0.
5.0
Figure 5: The normalized NTK map for 2-layer ReLU MLP with and without layerorm. The input and hidden dimensions are and 10,000. We can see that for the MLP without LayerNorm, prediction change 25 has dramatic influence on points far away like (4, meaning slight change of I (xo) will change (x) dramatically 4~.x.... However, MI.P equipped with layernorm exhibits good local property.
which suggests that even slight increase (or decrease) in the prediction value at To will result in an amplified increase (e.g.. x.5x) in the prediction value a distant point(e.g.. (4, 4)). Note that indeed the absolute NTK surrounding Zo is positive. However, values farther from are significantly larger. As result, the NTK around To is minimal and normalized predictions of the dataset sample and extreme points have large NTK value and exhibit a strange but strong correlation. As Q-value network is iteratively updated, it also tend to output large values for these extreme points, which makes extreme points easily become policy actions in optimization.
value close to zero. Thus, the value
This abnormal behavior contradicts the typical understanding of "kernel function" which usually diminishes with increasing distance, which represents the MLP's inherent limitation in accurate extrapolation. This indicates an intriguing yet relatively under-explored approach avoid divergence: regularizing the model's 's generalization on out-of-distribune an improperly large kernel value is that the neural network becomes a linear function when the input's norm is too large [54]. A more detailed explanation about linearity can be found at Appendix G. Therefore, simple method to accomplish this would be insert LayerNorm prior to each non- linear activation. We conduct a similar visualization by equipping the MLP with LayerNorm [3]. As shown Figure right, the value reaches its peak at æn and diminishes as the distance grows, demonstrating excellent local properties for well-structured kernels.
predictions. The main reason for such
Such architectural change beneficial for controlling the eigenvalue A by reducing the value of improperly large entries since matrix inequality tells us Amax (A) < || Therefore, it is expected that this method can yield small SEEM value, ensuring the convergence of training with minimal bias on the learned policy. We also provide theoretical justification explaining why LayerNorm results in lower SEEM value in the supplementary material. This explains why LayerNorm, as an empirical practice, can boost performance in previous online and offline RL studies [41; 5; 25; 28]. Our contribution is thus two-fold: We make the first theoretical explanation for how LayerNorm mitigates divergence through the NTK analysis above. 2) We conduct thorough experiments to empirically validate its effectiveness, as detailed Section Section 5.
10
Estimated Q-value
1.0
SEEM
Estimated Q-value
0.04
SEEM
10
0.9
0.03
0.6
0.02
0.4
10'
0.2
0.01
10"
0
250
500
750
1000
0.0
250
500
750
1000
250
500
750
1000
0.00
260
500
750
1000
steps
sleps
steps
sleps
Whight Norm
Weight
Dropout
Behavior Cloning
LayerNorm
BatchhNorm
Figure 6: The effect of various regularizations and normalizations on SEEM and Q-value. LayerNorm yields low SEEM and achieves stable Q-learning. Given the substantial disparity in the y-axis range among various regularizations, we present the results using two separate figures.
To validate effectiveness of LayerNorm practical offline RL settings, we select the walker2d- medium-expert-v2 task in D4RL to showcase how Q-value and SEEM evolves as the training proceeds. For comparative analysis, we also consider four popular regularization and normalization techniques, including BatchNorm [20], WeightNorm [44], dropout [45], and weight decay [35]. We
7

---

Table 1: Averaged normalized scores of the last ten consecutive checkpoints Antmaze tasks 10 seeds. diff-QL [51] reports the best score during the whole training process in its original paper. We rerun its official code for the average score. We also compare our method with diff-QL. by the best score metrics in the parenthesis. The best score of diff-QL is directly quote from its paper.
Dataset
TD3+BC
IQL
MSG
sIBC
diff-QL
ours
antmare-umaze-vO
40.2
87.5
98.6
93.3
95.6 (96.0)
94.3 + 0.5 (97.0)
antmaze- untmaze-umaze- umaze-divve diverse-v0
58.0
62.2
76.7
86.7
69.5 (84.0)
88.5 6.1
antmaze- medium- play--v-
0.2
71.2
83.0
88.3
(79.8)
85.6 1.7 (92.0)
0.0
70.0
83.0
90.0
6.2.0)
83.9 1.6 (90.7)
antmaze- cattmaze-larg-paee large-play-v0
39.6
46.8
63.3
(49.0)
6 (74.0)
antmaze-large-de large-diverse-v0
0.0
47.5
58.2
41.7
4.4
67.1 + 1.8 (75.7)
average
16.4
63.0
74.4
77.2
29.6 (75.4)
80.8 (87.4)
use TD3 as the baseline algorithms. To simplify our analysis, the target network is chosen to be the current Q-network with gradient stop instead of an Exponential Moving Average (EMA) of past networks. In our experiments, it is observed that LayerNorm and WeightNorm constrain SEEM and restrain divergence. However, weight decay and dropout did not yield similar results. As illustrated in Figure 6, weight decay and dropout evolve similar to the unregularized TD3 w.o. EMA baseline. WeightNNrm reduces SEEM by a clear margin compared to the baseline, thus demonstrating slighter divergence. In the right figure, we observe that behavior cloning effectively lowers SEEM at the cost of introducing significant explicit bias. Importantly, LayerNorm achieves both low SEEM and stable Q-values without necessitating explicit policy constraints. A notable outlier is BatchNorm. BatchNorm attains a relatively low maximum eigenvalue at the beginning of training but experiences an increase over time. Correspondingly, the Q-curve displays substantial oscillations. This instability could be ascribed to the shared batch normalization statistics between the value network and the target network, despite their input actions originating from distinct distributions [6].
Agent Performance
Previous state-of-thes ofete---rre--ttttt offline RL algorithms have performed exceptionally well on D4RL Mujoco Locomotion tasks, achieving an average score above 90 [2: 41]. In this section, we compare our method with various baselines two difficult settings, i.e., Antmaze and X% Mujoco task, to validate the effectiveness of LayerNorm in offline settings. Further experiments can be found at Appendix F.
5.1 Standard Antmaze Dataset
In Antmaze tasks characterized by sparse rewards and numerous suboptimal trajectories, the prior successful algorithms either relies on in-sample planning, such as weighted regression [26; 8], or requires careful adjustment the number of ensembles per game [15]. Algorithms based on TD3 SAC failed to achieve meaningful scores by simply incorporating a behavior cloning (BC) term [12]. Even Diff-QL, which replaces TD3+BC's Gaussian policy with expressive diffusion policies capture multi-modal behavior [51], continues to struggle with instability, leading to inferior performance.
We first show policy constraint (BC) is unable control q- value divergence while performing well in some challenging environments, by using Antmaze- large- play as an running example. We conduct experiments with Diff-QL by varying the BC (i.e., diffusion loss) coefficient from 0.5 to 10. As shown in Figure 7, when the policy constraint is weak (BC 0.5), it initially achieves a decent score, but as the degree of off-policy increases, the value starts diverge, and performance drops zero. Conversely, when the policy constraint is too strong (BC 10), the learned policy cannot navigate out of the maze due to suboptimal data, and performance remains zero. In contrast, simply incorporating LayerNorm into Diff-QL, our method ensures stable value convergence under less restrictive policy constraints (BC 0.5). This results in consistently stable performance in the challenging antmaze-large-play task. We report the overall results in Table 1, which shows that our method consistently maintains stable average scores across six distinct environments. Additionally, we are able to boost the highest score, while adopting the same evaluation metric used by Diff-QL. Ultimately, our method achieves an average score of 80.8 on Antmaze tasks, exceeding the performance previous SOTA methods. In summary, our experimental findings demonstrates within suboptimal datasets, an overly strong policy constraint is detrimental, while a weaker one may lead to value divergence. LayerNorm proves effective in maintaining stable value convergence under less restrictive policy constraints,

---

Estimated Q-value
Performance of Antmaze-large-play
BG coef cant
BC coefiiiinr
0.5
0.5
16
Juuns
Jumno
4
2
600
800
200
400
600
800
1000
K steps
K stepp
Figure 7: The detrimental bias in policy constraint. Although strong policy constraints succeed in depressing divergence, constraining the policy around sub-optmmll behavior sacrifices performance. With LayerNorm regularize SEEM without introducing a too strong bias, our method can achieve both stable Q-value convergence and excellent performance. The shallow curves represent the trajectory of individual seeds, while the darker curve denotes the aggregate across seeds.
resulting in an excellent performance. The ablation study regarding the specific configuration of adding LayerNorm can be found at Appendix F.4. The effect of other regularizations in Antmaze can be found at F.3. We also experimented with the regularizer proposed in DR3 [29]. Despite conducting thorough hyperparameter search for the coefficient, was observed that it had no substantive impact on enhancing the stability and performance of diff-QI..
5.2
X% Offline Dataset
Past popular offline RL algorithms primarily require transitions rather than full trajectories for training [40; 31; 26; 12; 51J. A superior algorithm should extrapolate the value of the next state and seamlessly stitch together transitions to form a coherent trajectory. However, commonly used benchmarks such as D4RL and RL Unplugged [10; 16] contain full trajectories. Even with random batch sampling, complete trajectories help suppress value divergence since the value of the next state 8141 will be directly updated when sampling {s1+1. a(+1, St+2, Pt-1) }. On the one hand, working with offline datasets consisting solely of transitions enables a more authentic evaluation of an algorithm's 's stability and generalization capabilities. On the other hand, particularly relevant in real-world applications such as healthcare and recommendation systems where it is often impossible to obtain complete patient histories.
We first evaluate the performance of offline algorithms within transition-based scenario. We construct transition- based datasets by randomly sampling varying proportions ()% from the D4RL Mujoco Locomotion datasets. Here, we set several levels for X E {1, 10, 50, 100}. When equals 100, equivalent D4RL dataset. As X decreases, the risk of encountering outtif-distiibution states and actions progressively esca- lates. To ensure fair comparison, we maintain the same subset across all algorithms. We evaluate every algorithm on Walker2d, hopper, and Halfcheetah tasks with medium, medium replay, and medium expert levels. For compari- son, we report the total score on nine Mujoco locomotion tasks and the standard deviation over 10 seeds. Figure 8 reveal marked drop in performance for all popular offline RL algorithms when the dataset is reduced to 10%. When the transition is very scarce (1 ()), all baselines achieve about only 50 to 150 total points on nine tasks. Further, we observed value divergence all algorithms, even for algorithms based policy constraints conservative value estimation.
800
T0D-BC
700
1OL
410L
800
training on the original
400
300
200
100
100%
50%
10%
1%
Figure 8: The performance of offline RL algorithms with the varying X% Mujoco Locomotion dataset. As the number of transitions decreases, all algorithms have a dramatic drop performance.
Subsequently, we demonstrate the effectiveness of LayerNorm improving the poor performance in X% datasets. By adding LayerNorm to the critic, value iteration for these algorithms becomes non- expansive, ultimately leading stable convergence. Figure depicts, under 10% and I 1 dataset, all baselines greatly improved by LayerNorm. For instance, under the 1 1 dataset, LayerNorm

---

enhances the performance of CQL from 50 points to 300 points, marking a 6x improvement. Similarly, TD3+BC improves from 70 228, representing approximately 3x improvement. This empirical evidence underscores the necessity of applying normalization to the function approximator for stabilizing value evaluation, particularly in sparse transition-based scenarios. Adding LayerNorm shown to be more effective than relying solely policy constraints. Another notable observation that with sufficient data (the 100% dataset) where previous SOTA algorithms
from Figure have performed excellently, LayerNorm only marginally alters the performance of algorithms. This suggests that LayerNorm could serve as a universally applicable plug-in method, effective regardless of whether the data is scarce plentiful.
TD8+BC
Ploffomanc Diffrrnne
200
150
100
80
100%
100%
100%
100%
Figure 9: The performance difference between baseline with LayerNorm and without it using the same X% dataset. The error bar represents the standard deviation over 10 seeds.
5.3 Online RL Experiments LayerNorm Allows Online Methods without EMA.
It is natural to ask whether our analysis and solution are also applicable to online settings. Previous works [34; 13; 17; 61] has empirically theoretically shown the effect of EMA stabilizing Q-value and prevent divergence in online setting. To answer the question, We tested whether LayerNorm can replace EMA prevent value divergence. We conducted experiments two environments in online settings: Hopper and Walker. Please refer to Figure 11 for the curves of return. Surprisingly, we discovered that LayerNorm solution allows the SAC without EMA to perform equivalently well as the SAC with EMA. In contrast, SAC without EMA and LayerNorm behaved aimlessly, maintaining near-zero scores. reveals LayerNorm (or further regularizations discovered by SEEM later) have the potential of allowing online DQN-style methods to step away from EMA/frozen target approaches altogether towards using the same online and target networks. These results partially reveals potential implications of SEEM in an online learning setting.
6 Conclusion and Limitations
In this paper, we delve into the Q-value divergence phenomenon offline RL and gave com- prehensive theoretical analysis and explanation for it. We identified a fundamental process called self-excitation as the trigger for divergence, and propose an eigenvalue measure called SEEM to reliably detect and predict the divergence. Based SEEM, we proposed an orthogonal perspective other than policy constraint to avoid divergence, by using LayerNorm to regularize the generalization of the MLP neural network. We demonstrated empirical and theoretical evidence that method introduces less bias in learned policy, which has better performance in various settings. Moreover, the SEEM metric can serve as an indicator, guiding future works toward further engineering that may yield adjustments with even better performance than the simple LayerNorm. Despite the promising results, our study has some limitations. we have not incorporated the impact of Exponential Moving Average (EMA) into our analysis. Also, the existence of critical point is still not well-understood. Moreover, our analysis does extend to online RL, as the introduction of new transitions disrupts the stability of the NTK direction. All these aspects are left for future work.
7 Acknowledgement
This work is supported by the National Key R&D Program China (2022D001149))
10

---

References
[1] Joshua Achiam, Ethan Knight, and Pieter Abbeel. Towards characterizing divergence in deep q-learning arXiv preprint wiXiv:/933008000 2019. 14
[2] Gaon An, Seungyong Moon, Jang-Hyun Kim, and Hyun Oh Song. Uncertainty-basel offline reinforcement learning with diversified q-ensemble. NIPS, 2021. 2,8
[3] Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E Hinton. Layer normalization. arXiv preprint arXiv:1607.064500 2016." 7
[4] Leemon Baird. Residual algorithms: Reinforcement learning with function approximation. In Machine Learning Proceedings. 1995. 1, 2, 14
[5] Philip Ball, Laura Smith, Ilya Kostrikov, and Sergey Levine. Efficient online reinforcement
learning with offline data. In ICML, 2023. 7, 14
[6] Aditya Bhatt, Max Argus, Artemij Amiranashvili, and Thomas Brox. Crossnorm: Normalization for off-policy td reinforcement learning. arXiv preprint arXiv:1902.05605, 2019. 8, 14
[7] David Brandfonbrener, Will Whitney, Rajesh Ranganath, and Joan Bruna. Offline without off-policy evaluation. NIPS, 2021. 9, 14
[8] Huayu Chen, Cheng Lu, Chengyang Ying, Hang Su, and Jun Zhu. Offline reinforcement learning via high-fidelity generative behavior modeling, In ICLR, 2023. 8, 14
[9] Lili Chen, Kevin Lu, Aravind Rajeswaran, Kimin Lee, Aditya Grover, Misha Laskin, Pieter Abbeel, Aravind Srinivas, and Igor Mordatch. Decision transformer: Reinforcement learning via sequence modeling. NIPS, 2021. 14
[10] Justin Fu, Aviral Kumar, Ofrr Nachum, George Tucker, and Sergey Levine. D4rl: Datasets for deep data-driven reinforcement learning, 2020. 4,
[11] Justin Fu, Aviral Kumar, Matthew Soh, and Sergey Levine. Diagnosing bottlenecks in deep q-learning algorithms. In /CMI, 2019. 14
[12] Scott Fujimoto and Shixiang Shane Gu. A minimalist approach to offline reinforcement learning. NIPS, 2021. 2, 8, 9,
[13] Scott Fujimoto, Herke Hoof, and David Meger. Addressing function approximation error in actor-critic methods. ICML, 2018. 10, 14
[14] Scott Fujimoto, David Meger, and Doina Precup. Off-policyy deep reinforcement learning without exploration. In ICML, 2019. 1, 14
[15] Kamyar Ghasemipour, Shixiang Shane Gu, and Ofir Nachum. Why so pessimistic? estimating uncertainties for offline rl through ensembles, and why their independence matters. NIPS, 2022.
[16] Caglar Gulcehre, Ziyu Wang, Alexander Novikov, Thomas Paine, Sergio Gómez, Konrad Zolna, Rishabh Agarwal, Josh S Merel, Daniel J Mankowitz, Cosmin Paduraru, et al. RI unplugged: A suite of benchmarks for offline reinforcement learning. NIPS, 2020. 9
[17] Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with a stochastic actor. In ICML, 2018. 1, 10
[18] Yizeng Han, Yifan Pu, Zihang Lai, Chaofei Wang, Shiji Song, Junfeng Cao, Wenhui Huang. Chao Deng, and Gao Huang. Learning to weight samples for dynamic early-exiting networks. In ECCV, 2022. 22
[19] Hado Hasselt. Double q-learning. NIPS, 2010. 1, 14, 23
[20] Sergey Ioffe and Christian Szegedy. Batch normalization: Accelerating deep network training by reducing internal covariate shift. In ICML, 2015. 7, 14

---

[21] Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural tangent kernel: Convergence and generalization neural networks. NIPS, 2018. 3,4
[22] Michael Janner, Qiyang Li, and Sergey Levine. Offline reinforcement learning as one big sequence modeling problem. In NIPS, 2021. 14
[23] Natasha Jaques, Asma Ghandcharioun, Judy Hanwen Shen, Craig Ferguson, Ágata Lapedriza. Noah Jones, Shixiang Gu, and Rosalind W. Picard. Way off-policy batch deep reinforcement learning of implicit human preferences in dialog. CoRR, 2019. 2
[24] Bingyi Kang, Xiao Ma, Chao Du, Tianyu Pang, and Shuicheng Yan. Efficient diffusion policies for offline reinforcement learning. In NIPS, 2023. 22
[25] Bingyi Kang, Xiao Ma, Yirui Wang, Yang Yue, and Shuicheng Yan. Improving and benchmark- ing offline reinforcement learning algorithms, 2023. 7, 14, 22
[26] Ilya Kostrikov, Ashvin Nair, and Sergey Levine. Offline reinforcement learning with implicit q-learning. In NIPS Deep RL Workshop, 2021. 8, 9, 14
[27] Alex Krizhevsky, llya Sutskever, and Geoffrey E Hinton. Imagenet classification with deep convolutional neural networks. NIPS, 2012. 15
[28] Aviral Kumar, Rishabh Agarwal, Xinyang Geng, George Tucker, and Sergey Levine. Offline q-learning on diverse multi-task data both scales and generalizes. In ICLR, 2023. 7, 14
[29] Aviral Kumar, Rishabh Agarwal, Tengyu Ma, Aaron Courville, George Tucker, and Sergey Levine. Dr3: Value-based deep reinforcement learning requires explicit regularization. In ICLR, 2022. 6, 9
[30] Aviral Kumar, Justin Fu, Matthew Soh, George Tucker, and Sergey Levine. Stabilizing off- q-learning via bootstrapping error reduction. NIPS, 2019. 1,2, 14
[31] Aviral Kumar, Aurick Zhou, George Tucker, and Sergey Levine. Conservative q-learning for offline reinforcement learning. In NIPS, 2020. 2, 9, 14
[32] Sascha Lange, Thomas Gabel, and Martin Riedmiller. Batch reinforcement learning. In Reinforcement learning. Springer, 2012. 14
[33] Sergey Levine, Aviral Kumar, George Tucker, and Justin Fu. Offline reinforcement learning: Tutorial, review, and perspectives on open problems. arXiv preprint arXiv:2005.01643, 2020.
[34] Timothy P. Lillicrap, Jonathan J. Hunt, Alexander Pritzel, Nicolas Heess, Tom Erez, Yuval Tassa, David Silver, and Daan Wierstra. Continuous control with deep reinforcement learning. In ICLR, 2016. 10
[35] Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. In ICLR, 2019. 7
[36] Rui Lu, Gao Huang, and Simon S. Du. On the power of multitask representation learning in linear mdp, 2021. 14
[37] Rui Lu, Andrew Zhao, Simon S. Du, and Gao Huang. Provable general function class represen- tation learning in multitask bandits and mdps. In NIPS, 2022. 14
[38] Xiao Ma, Bingyi Kang, Zhongwen Xu, Min Lin, and Shuicheng Yan. Mutual information regularized offline reinforcement learning. In NIPS, 2023. 2, 14
[39] Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Andrei A Rusu, Joel Veness, Mare G Bellemare, Alex Graves, Martin Riedmiller, Andreas K Fidjeland, Georg Ostrovski, et al. Human-level control through deep reinforcement learning. nature, 2015. 1, 14
[40] Ashvin Nair, Abhishek Gupta, Murtaza Dalal, and Sergey Levine Awac: Accelerating online reinforcement learning with offline datasets. arXiv preprint arXiv:2006.099: 09359, 2020. 9, 14
[41] Alexander Nikulin, Vladislav Kurenkov, Denis Tarasov, Dmitry Akimov, and Sergey Kolesnikov. Q-ensemble for offline Don't scale the ensemble, scale the batch size. arXiv preprint arXiv:2211.11092 2022. 2, 7, 8, 14
12

---

[42] Xue Bin Peng, Aviral Kumar, Grace Zhang, and Sergey Levine. Advantage- regression: Simple and scalable off-po reinforcement learning. arXiv preprint arXiv:1910.00177, 2019.
2, 14
[43] Yifan Pu, Yiru Wang, Zhuofan Xia, Yizeng Han, Yulin Wang, Weihao Gan, Zidong Wang, Shiji Song, and Gao Huang. Adaptive rotated convolution for rotated object detection. In ICCV, 2023. 15
[44] Tim Salimans and Durk P Kingma. Weight normalization: A simple reparameterization to accelerate training of deep neural networks. NIPS, 2016. 7
[45] Nitish Srivastava, Geoffrey Hinton, Alex Krizhevsky, Ilya Sutskever, and Ruslan Salakhutdinov. Dropout: a simple way to prevent neural networks from overfitting. JMLR, 15(1):1929-1958 2014.7
[46] Richard S Sutton and Andrew G Barto. Reinforcement learning: An introduction. MIT press, 2018. 1, 14
[47] JN Tsitsiklis and B Van Roy. An analysis of cemporal-differnnee learning with function approximationtechnical Rep. LIDS-P-2322) Lab. Inf. Decis. Syst. Massachusetts Inst. Technol. Tech. Rep. 1996. 1, 2, 14, 17
[48] Hado Van Hasselt, Yotam Doron, Florian Strub, Matteo Hessel, Nicolas Sonnerat, and Joseph Modayil. Deep reinforcement learning and the deadly triad. arXiv preprint arXiv:1812.0266 0264822488. 2018. 1, 14, 23
[49] Chaofei Wang. Jiayu Xiao, Yizeng Han, Qisen Yang, Shiji Song, and Gao Huang. Towards learning spatially discriminative feature representations. In ICCV, 2021. 15
[50] Chaofei Wang, Qisen Yang, Rui Huang, Shiji Song, and Gao Huang. Efficient knowledge
distillation from model checkpoints. NIPS, 2022. 15
[51] Zhendong Wang, Jonathan J Hunt, and Mingyuan Zhou. Diffusion policies expressive policy class for offiine reinforcement learning. In ICLR, 2023. 2, 8, 9, 14, 23
[52] Ziyu Wang, Alexander Novikov, Konrad Zolna, Josh Merel, Jost Tobias Springenberg, Scott Reed, Bobak Shahriari, Noah Y. Siegel, Çaglar Gülçehre, Nicolas Heess, and Nando de Freitas. Critic regularized regression. In NIPS, 2020. 14
[53] Yifan Wu, George Tucker, and Ofir Nachum. Behavior regularized offline reinforcement learning. arXiv preprint arXiv: /9/1. 19361, 2019. 2, 14
[54] Keyulu Xu, Mozhi Zhang, Jingling Li, Simon Shaolei Du, Ken-Ichi Kawarabayashi, and Stefanie Jegelka. How neural networks extrapolate: From feedforward to graph neural networks. In ICLR, 2021. 7, 25
[55] Qisen Yang, Shenzhi Wang, Matthicu Gaetan Lin, Shiji Song, and Gao Huang. Boosting offline reinforcement learning with action preference query. ICMI., 2023. 14
[56] Qisen Yang, Shenzhi Wang, Qihang Zhang, Gao Huang, and Shiji Song. Hundreds guide millions: Adaptive offline reinforcement learning with expert guidance. arXiv preprint arXiv:2309.014.4 2023. 14
[57] Tianhe Yu, Aviral Kumar, Rafael Rafailov, Aravind Rajeswaran, Sergey Levine and Chelsea
Finn. Combo: Conservative offline model-based policy optimization. NIPS, 2021. 2
[58] Yang Yue, Bingyi Kang, Xiao Ma, Gao Huang, Shiji Song, and Shuicheng Yan. Offline prioritized experience replay, arXiv preprint arXiv:2306.0544 32230000.122. 12, 2023. 22
[59] Yang Yue. Bingyi Kang, Xiao Ma, Zhongwen Xu, Gao Huang, and Shuicheng Yan. Boosting offline reinforcement learning via data rebalancing. In Offline RI. Workshop. 2022. 22
[60] Yang Yue, Bingyi Kang, Zhongwen Xu, Gao Huang, and Shuicheng Yan. Value-consistent representation learning for data- reinforcement learning. In AAAI, 2023. 14
[61] Shangtong Zhang, Hengshuai Yao, and Shimon Whiteson. Breaking the deadly triad with target network. In ICML, 2021. 10
13

---

A Related Works
Value Divergence with Neural Network. In online reinforcement learning (RL.), off- algo- rithms that employ value function approximation and bootstrapping can experience value divergence, phenomenon known the deadly triad [46; 4; 47; 48; 11]. Deep Q-Networks (DQN) typify this issue. As they employ neural networks for function approximation, they are particularly susceptible to Q-value divergence [19; 13; 48; 60]. Past research has sought to empirically address this divergence problem through various methods, such as the use of separate target networks [39] and Double-Q learning [19; 131. Achiam et al. [1] analyze linear approximation of Q-network to characterizes the diverge, while CrossNorm [6] uses a elaborated version of BatchNorm [20] to achieve stable learning. Value divergence becomes even more salient in offline RL, where the algorithm learns purely from a fixed dataset without additional environment interaction [14; 30j. Much of the fo- cus in the field of offline RL has been on controlling the extent of off- -olppollyy learning, i.e. policy constraint [12; 40; 26; 14; 53; 51; 8; 42]. Several previous studies [41; 5; 25; 28] have empirically utilized LayerNorm to enhance performance in online and offline RL. These empirical results partially align with the experimental section of our work. However, our study makes theoretical explanation for how LayerNorm mitigates divergence through the NTK analysis. Specifically, we empirically and theoretically illustrate how LayerNorm reduces SEEM. In addition to LayerNorm, our contri- bution extends to explaining divergence and proposing promising solutions from the perspective of reducing SEEM. Specially, discover that WeightNorm can also be an effective tool and explain why other regularization techniques fall short. Finally, we perform comprehensive experiments to empirically verify the effectiveness of LayerNorm the %X dataset, a practical setting not explored in previous work. Thus, our contributions are multifaceted and extend beyond the mere application of LayerNorm.
Offline RL. Traditional RL and bandit methods have a severe sample efficiency problem [32; 14; 36; 37]: offline RL aims to tackle this issue. However, offline RL presents significant challenges due to severe off fff-policy issues and extrapolation errors. Some existing methods focuses on designs explicit or implicit policy regularizations to minimize the discrepancy between the learned and behavior policies. For example, TD3+BC [13; 12; 55; 56] directly adds a behavior cloning loss to mimic the behavior policy, Diffusion-QU QL [51] further replace the BC loss with a diffusion loss and using diffusion models as the policy. AtR-style [42; 52; 7] and IQL [26] impose an implicit policy regularization by maximizing the weighted in-distribution action likelihood. Meanwhile, some other works try to alleviate the extrapolation errors by modifying the policy evaluation procedure. Specifically, CQL [31] penalizes out-of distribubtii actions for having higher Q-values, while MISA [38] regularizes both policy improvement and evaluation with the mutual information between state-action pairs of the dataset. Alternatively, decision transformer (DT) [9] and trajectory transformer [22] cast offline RL as a sequence generation problem, which are beyond the scope of this paper. Despite the effectiveness of the above methods. they usually neglect the effect of function approximator and are thus sensitive to hyperparameters for trading off performance and training stability. Exploration into the function approximation aspect of the deadly triad is lacking in offline RL. Moreover, a theoretical analysis of divergence in offline RL that does not consider the function approximator would be inherently incomplete. We instead focus on this orthogonal perspective and provide both theoretical understanding and empirical solution to the offline RL problem.
B Proof of Main Theorems
Before proving our main theorem, we first state an important lemma.
Lemma 1. For L-layer ReLU-activate MLP and fixed input x. x', assume the biases are relatively negligible compared to the values in the hidden layers. If we scale up every parameter of fo to X times, namely 0' = 20, then we have following equations hold for any x, x' S .Sxx A
[a) - X2 fo =o(1) (x))l)
Vefo - T. Vo.e((=) =o(1) (
ko (z, x²) X2- 2 X2(L-11 ko x0)
=o(1) ka (a. x').
Proof. Recursively define
* Wi + be, 20 = x
14

---

Σ = o(t)).
=
Then it is easy to see that if we multiply each We e and by by X, denote the new corresponding value to be z/c and note that the bias term is negligible, we have
= =
2/2 =XW 22 + Ab2
=X2(W + +b2) (A 2R)b2
=X²22 (A- -22
=X²22 + 62
(negilible bias ||22|| = o(19))
2/11 (2)) + Abet1
=XWio(X'z) ce) Abe+1
= 10 (z(+11) Abb+1 + AW gd;
(6 o('2; + cx) o(2(411))
= + (A - Attlybe+1 + AW 55
Note that dell A||
= 2/41111- the last step is because (-)) is I-Lipshitz, and we have (X' zς + =t+1 = Eş and ||e/|| o(X)) When value divergence is about to occur, the dot product W,,, becomes relatively large compared to the scalar bi and the effect of bias on the output becomes negligible. Therefore, we can recursively deduce that the remaining error term
+c) 0( > W 11
x + 2wede
is always negiligle compared with 2/441- Hence we know for (zx) is close to fox) with only o11) relative error.
Taking gradient backwards, we know that fore we know
0f aw,
is proportional to both || and 81
There-
J=
0.e 0.0: 02'411
0.e 2 0fe 0zt+1
2/
X²g, x
aw;
oW. 6
This suggests that all gradients with respect to the weights become scaled by factor of X The gradients with respect to the biases are proportional X When X is large enough to render the gradient of the bias term negligible, it follows that |e (z) - VVoooo) o(1) (()) e (ar) This equation implies that the gradient updates for the model parameters are dominated by the weights, with negligible contribution from the bias terms. And since NTK the inner product between gradients, we know Gø (æx, æ²''s relative error to X2(L-1) Go( x!) is (1) negligible. We run experiments in Appendix F to validate the homogeneity. We leave analysis for Convolution Network [27] and its variants [43; 49; 50] for future work.
Theorem Suppose that the network's parameter al iteration t is O,, For each transition (st, Gi Tx) in dataset, denote r = [ri..... ........Ta. .T]] T RM, (t. = arg max, (s, Denote of, ($141, fe, (8t+1)). Concatenate all x;,, to be X Denote u, fe. (X) (r +y fo, (X))) to be TD error vector at iteration t. The learning rate rj is infinitesimal. When the maximal point of Qo. is stable as t increases, we have the following evolving equation for u4+1
U (I + +1.
where A = (700, (Xi) dh, (X)) 60. (X) Goo (X; X) (X,
Proof. For the sake of simplicity, denote Z, = Vofo(X) Z = vefe(X)) The Q-value
iteration minimizes loss function L defined by L(0) = llo(( - (r+ x So, (x; Therefore
we have the gradient as
OC(0) = Z: (fo(X) - (r + fo. (xi)) = Z,4.
(4)
00
15

---

According to gradient descent, we know Bet1 = 0; - Since is very small, we know 0 stays within the neighborhood of 0,. We can Taylor-expand function f.) near 0, as
fo(X) ( (0- 0,) Jo. (X) 2 00 0+) + (X).
(5)
(x 2 (Z)) (0- 0.) + So, (x
(6)
When is infinitesimally small, the equation holds. Plug in 0c+1, we know
fert1 (X) fe. (X) = -, z, = 17 Go. (X, X), S0..1 (X?) fo. (X;) (z,) Z;u, -1 Ge. (X;, X ju:-
(7)
(8)
According to Assumption 2, we know ||x; - x+11l = = So we can substitute X;+1 into X; since it is high order infinitesimal. Therefore, U1+1 boils down to
uc+1 for.1 (X) r - (Xi+1)
(9)
(X.() 17- Go. (X, x), - >(fe, -17 11. Ge,((;; + [()))
(10)
(X) - r rfo,
tn (X;; (X; G X) + o(1)
)
(11)
ui
= +I+ =A+nAi))u
(12)
where A = (Xi,, (X,
Theorem 3 Suppose we use SGD optimizer for Q-value iteration with learning rate 1j to be infinitesi- G X X), where G is the Gram matrix under
mal. Given iteration > lo, and = -)- terminal kernel k. The divergence of u, is equivalent to whether there exists an eigenvalue À of A such that Re(A) > 0. If converge, we have = (I +++,44 nt) Uta Otherwise, u, becomes parallel to the eigenvector of the largest eigenvalue X of A, and its norm diverges to infinity at following order
| 0 C1 (' (1
(13)
for some constant C' to be determined and L is the number of layers of MLP. Specially, when L 2, it reduces to O x-
Proof. According to Assumption 2, max action becomes stable after to- It implies IX; -x--x" Xx1 o().
The stability of the NTK direction implies that for some scalar h; and the specific input X X, we have Ge. (X", x)= Gx"" X) + E,11 and Go. (X. GGG, X) + c1 where ||e ||
o(k,), = 2, Therefore, we have
A. G X) Go, (X =,G(Xi, - G(X, + (EI,1 E1,2)
(14)
(15)
=k(GG(XX' ( -X)- (7E1,3 + Et,1 - 61,2)
(16)
(E,.3 = G(X, -
X))
= 8.
(17)
k2 equals if the training is convergent, but will float up if the model's predicted Q-value blows up. According to Assumption 2, the rest term ö, = 7E,,3 + Et,1 E, is of o(1) norm and thus negligible, we know all the eigenvalues of I + 1 have form 1 nd2 Considering 1) small enough, we have [1 +1++1 22, ~ 1 + 2n Re(A). Now suppose if there does not exists eigenvalue À of A satisfies Re(A) > 0, we have [1 + ./111 S 1. Therefore, the NTK will become perfectly stable so k; = 1 for t > to. and we have
u. = -
7A AA,,)), -...- (I
A.)uuo n
(18)
1-1
= II (I + mA)u: = (A) Ukg
(19)
==1
16

---

Otherwise, there exists an eigenvalue for A satisfying Re(A) > O. Denote the one with the largest real part as and be the corresponding cigenvector. We know matrix I + n/A also has left eigenvector v, whose eigenvalue is v/A. In this situation, we know after each iteration, || || will become larger than |u:]l- Moreover, to achieve larger and larger prediction values, the model's parameter's s norm ||O||| also starts to explode. We know the is homogeneous with respect 0, for ReLU networks. The output fo, (X) enlarges p² times when 0, enlarges p times. When the reward values is small with respect to the divergent Q-value, TD error ue = ffe (X)) = 0(0)). Besides, according to Lemmal, we know 4 = [[ 12((--1) = (1] [/ || 1)/1) =
2/4).
Uz. left multiply v to equation 1 (I + A we have g(nt + 11) 11+k)))gg((1 Since we know such iteration will let ue to be dominated by and align with v, we know g(nl) = O 0[1|uell for large f. Therefore k = O[1
Denote ((n))
= C g(11)2-2/2 This boils /|2(L-1)/2) ([]],,,,(((_)))))
down ggp + 77) g(nt) + CrdAg(nt)2 which further becomes
g(nl + 1) g(rl = ('g(n1t)3221
(20)
Let 1) - O, we have an differential equation dé degenerates to linear function. The solution of ODE is
C2Ag(t)3-274 2//3 When 1, the MLP network
| = g(rt) = C'e exe
(21)
reflecting the exponential growth under linear function that has been studied in previous works [47]. When I > 2, Solving this ODE gives
1
g(t) (1 (1- C-A) C'AL) 2/(22-2)
(22)
So at an infinite limit, we know |/u/||| = g(mit) = 0 (1-C"-n)--<aa- Specially, for the exper-
imental case we study in Figure 2 where L = 2, it reduces to O ( We conduct more experiments with L 3 in Appendix C.2 to verify our theoretical findings.
also worthwhile to point out another important mechanism that we found when the Q-value network is in steady state of diverging. Although we assume that the learning rate 1) is infinitesimal, the arg max function is not L-iishhitz So why could both NTK function and optimal policy action X; converge in direction? We found that when the model enters such "terminal direction", (ss, fto (8:) not only becomes the max value state-action pair, but also satisfies
too (s:) arg max AEA
Ke, (Za: 24), Ea = (ss, a).
This means each ,;,, = (s,, (s:)) also gain the maximum value increment among all the possible actions, maintaining its policy action property. Such stationary policy action for all state causes the model's 's parameter to evolve in static dynamics, which in turn reassures it to be policy action in the next iteration. Therefore, after the critical point to, model's parameter evovlves linearly to infinity, its NTK and policy action becomes stationary, and all these factors contribute to the persistence of each other.
C More Observations and Deduction
C.1 Model Alignment
In addition to the findings presented in Theorem 1 and Theorem 3, we have noticed several intriguing phenomena. Notably, beyond the critical point, gradients tend to align along particular direction, leading to an infinite growth of the model's parameters that same direction. This phenomenon supported by the observations presented in Figure 16, Figure 17, and Figure 18, where the cosine similarity between the current model parameters and the ones at the ending of training remains close to I after reaching critical point, even as the norm of the parameters continually increases.
17

---

C.2
Terminal Time
Theorem claims ||u+|| = -
=0(
-a((-oon))rea-) implying the relation
1/9(22-2)/ 2114 1 -1
(23)
Beisdes, it implies the existence of "terminal time" Č" that the model must crash at singular point. When the training approaches this singular point, the estimation value and the model's 's norm explode rapidly in very steps. We have run an experiment with L = 2 in Figure from which we can see that Q-value's inverse proves to decay linearly and eventually becomes Nan at the designated time step. When I, 3, from our theretical analysis, we have 1/q² 1 C" Xny.. The experimental results in Figure 10 corroborate this theoretical prediction, where the inverse Q-value raised to the power of 4/3 is proportional to C" Ant after critical point and it eventually reaches NAN value at the terminal time step.
1/q 4/3 V.s. step
3,0
R-squared = 0.99855
2.5
y = 0.00270 x 2.83
2.0
1.5
1.0
0.5
0.0
400
600
1000
Steps
Figure 10: Linear decay with SGD and L=3,
C.3 Adam Case
In this section, will prove that if the algorithm employs Adam the optimizer, the model still suffers divergence. Moreover, we demonstrate that the norm of the network increase linearly, of which the slope is 7VP, where P is the number of parameters and 17 is the learning rate. Also, the Q-value prediction will increase L pollyoooommi 's rate, where L is the number of layers of model fo. Experimental results in Figure 4 verified our findings. Besides, we show that all runnings across D4RL. environments represents the linear growth of the norm of the Q-network Figure 16, Figure 17, and Figure 18.
Theorem 4. Suppose
Adam optimizer for Q-value iteration and all other settings the same as Theorem 3. After t > to, the model will diverge if and only if >mxx > 0. If it diverges, we have ||0,]| nVtPt o(t) and ||well O(t²) where P and L are the number of parameters and the number of layers for network fo. respectively.
Proof. We only focus on the asymptotic behavior of Adam. So we only care about the dynamics for t T for some large T. Also, at this regime, we know that the gradient has greatly aligned with the model parameters. So we assume that
VL(0,) = -C, C>0
(24)
18

---

Recall that each iteration of the Adam algorithm has the following steps.
9t VI((O);
(25)
= -,m,-1 + (1-Bi)9:.
(26)
01 82 + (1 (1-32)98
(27)
mg
m 1 81
(28)
Ut
6 = 1 82
(29)
17
0,41 = 0g V0t
(30)
Instead of exactly solving this series, we can verify linear growth is indeed the terminal behavior of O, since we only care about asymptotic order. Assume that 0, kt for > T,, we can calculate m, by dividing both sides of the definition m, by Bi, which gives
m
my
3
B
B
St-
(31)
m4
31
Σ
1
89
31
gs.
(32)
==0
my
C
B1), (1-13)88
-
(33)
where Ir is given in Equation (24). Similarly, we have
Ut = k0 + o(t
(34)
Hence we verify that
m4
1
3 Ut
kCt
0c+1 0:
17
11
CP2 =
3
therefore we know each iteration will increase each parameter by exactly constant 1). This in turn verified our assumption that parameter 0, grows linearly. The slope for the overall parameter is thus 7VP. This can also be verified in Figure 4. When we have 0, 7//6 where is the normalized parameter, we can further deduce the increasing order of the model's estimation. According to lemma 1. the Q-value estimation (also the training error) increase at speed (
D LayerNorm's Effect on NTK
In this section, we demonstrate the effect LayerNorm on SEEM. Our demonstration just an intuitive explanation rather than rigorous proof. We show that adding LayerNorm can effectively reduce the NTK between any To and extreme input x down from linear to constant. Since each entry of Gram matrix G is an individual NTK value, we can informally expect that G(X; X)'s 's cigenvalue are greatly reduced when every individual NTK value between any xy and extreme input is reduced.
We consider two-layer MLP. The input x E Rdd,, and the hidden dimension d. The parameters include W [w,,... wa)" E Igoxden b E Rd and E Rd. Since for the NTK value, the last layer's 's bias term has a constant gradient, we do not need to consider it. The forward function of the network is
fo(x) Σ (w) x bz).
Proposition 1. For any input and network parameter 0, if Vefo(x) + O, then we have
lim kNTK (x, Ax) = S((A) - ac.
(35)
X
19

---

Proof. Denote 2, w; æ + b, according to condition Vefo(x) # O, we know there must exist at
least one such that 2j > denote this set as P. Now consider all the [d] that satisfy 2, > 0 and
w; (otherwise take opposite sign of X), have
Of
ow x + b.) =w, x + bi.
(36)
0a;
0f
=0
(37)
Ow;
m
0f
-(+-
(38)
0b. m
Similarly, have
0f
(Aw, b.) xeo b..
(39)
0a, Ax
(40)
Omn
Axe
0s =0,
(41)
σb, Aa
So have
0f(x)
(Aar
Σ RE
= > ((w, x)² + b. + a + = O())
00;
00,
Denote
N {1, (l....... d} P. We know for every N either 0((x) 0f(2)
0f(w)
Ow., 0, or
80, 8b,
w) x < 0. For the latter case, we know lima-100 have
0f(Am)
0f(Az)
0f(Am)
0. In both cases, we
da;
ab;
ow:
0f(x)
0f(xx)
lim A-000
= 0.
00,
00;
JEN
Therefore, according to the definition of NTK, we have
0f(z)
0f(xx) O(A) co.
lim KNTK (z, Aas)
lim
X
X
00;
00;
For the model equipped with LayerNorm, the forward function becomes fo(x) ofo(Wx b)),
where (-) is the layer normalization function defined as
x 11 x/d
(x) vd
|| 11 T a/d]]
Denote P=I-11 /d, note that the derivative of vb(-) is
Vd. I
Pax P
((x) dx
P.
(42)
|Pol
P2x 3
Specially, we have
λx X11 x/d
((x)) Vd-
((x)..
(43)
x 11 x/d|
Now we state the second proposition.
Proposition 2. For any imput and network parameter and any direction E Rddi,, if the network has LayerNorm, then we know there exists a universal constant C, such that for any X 2 0, we have
kYTK + Xx) S C.
(44)
20

---

Proof. Since for finite range, there always exists constant upper bound, we just need to analyze the case for + +00 and shows that is constant bounded. First compute Vofo and get
0f
(W
(45)
da m
0ʃ
=a
o'( =a
b)) b)æ,
(46)
OW E
0s =a =a =aTo'(o(Wa +)) (W + b).
(47)
0b æ
These quantities are constant bounded. Next we compute lim Vofo(æ Xe)
0
(x Am) b))); (Wb(WW
(48)
da
=a caTo( b)) λv) + 6) x + Av),
(49)
OW
m Av
Of
=a
(a) + b).
(50)
0b
According to the property of LayerNorm in Equation (43), we have
0f
lim A
lim (W + (w) + b))
(51)
Da a+
A0000
=o(b(W (W (Av))) =o(b(W v)) Constant
(52)
(53)
aʃ
Tim
lim a T o' + Xo) + bW) + Xv) + b) Av)
(54)
X-000
OW a+
lim a T (W)()( W ( + Av) + b) + Xv)
(55)
X=00
I
P(AW (A) P P((We)]
lim TT (' X
P(x Xv)
PAWe
(56)
P(x + Xv) PW
PWev T w P(x + Av)
lim a (We))V X
PW
(57)
-a T (WWWe))))
Pe |Wool
PW W Pv
(58)
PW
Constant.
(59)
lim 0f
0b x+
lim a (oWo)(V
P
I
PW v) P
P
P
(60)
-0.
(61)
Therefore we know its limit is also constant bounded.
So we know there exists a universal constant (0f(m) 0f(a-
with respect to 0, x, v such that KNTK (x: x + Au)) = 00:
I C.
00,
E Experiment Setup
SEEM Experiments For the experiments presented in Section Section 3.1, we adopted TD3 as our baseline, but with a modification: instead of using an exponential moving average (EMA), we directly copied the current Q-network as the target network. The Adam optimizer was used with a learning rate of 0.0003, 81 = 0.9, and 32 0.999. The discount factor, Y, was set to 0.99. Our code builds upon the existing TD3+BC framework, which can be found at https /ttgt/gttt//////// con/ com/sfujin sfujim/TDD BC.
21

---

SEEM Reduction Experiments For the experiments discussed in Section Section 4, we maintained the same configuration as the SEEM experiments, with the exception of adding regularizations and normalizations. LayerNorm was implemented between the linear and activation layers with learnable affine parameters, applied to all hidden layers excluding the output layer. WeightNorm was applied to the output layer weights.
Offiine RL Algorithm Experiments For the X% Mujoco Locomotion offline dataset experiments presented in Section Section 5, we used true offline RL algorithms including TD3+BC, IQL., Diff- QL. and CQL as baselines. We implement our method on the top of official implementations of TD3+BC and IQL; for CQL,
reliable JAX implementation at https: : //gg/gtt thub com/ sail-sgooffbench [25]; for Diffusion Q-learning (i.e., combining TD3+BC and diffusion policy), efficient JAX implementation EDP [24]. LayerNorm was directly added to the critic network in these experiments. For the standard Antmaze experiment, we assign a reward +10 to the successful transitions, otherwise O. Considering the extreme suboptimality in Antmaze medium and large datasets, we adopt Return- Offline Prioritized Replay (OPER-R) R) with Phace = 0 to rebalance these suboptimal datasets [58; 59; 18]. Specifically, instead of uniform sampling during policy improvement, we assign a higher probability to the transitions in the successful trajectories. Meanwhile, a uniform sampler is still used for policy evaluation. When removing OPER-R, diff-QL. with LayerNorm can still achieve a stable performance with slightly lower scores (see Table 2).
Linear Decay of Inverse Q-value with SGD Given that the explosion in D4RL environments occurs very quickly in the order of x-j and is difficult to capture, we opted to use a simple toy task for these experiments. The task includes a continuous wwo-dimensional state space s = (x1, .22 E = R2, where the agent can freely navigate the plane. The action space is discrete, with 8 possible actions representing combinations of forward or backward movement in two directions. Each action changes the state by value of 0.01. All rewards are set to zero, meaning that the true Q-value should be zero for all state-action pairs. For this task, randomly sampled 100 state-action pairs as our offline dataset. The Q-network was implemented as two-layer MLP with hidden size of 200. We used SGD with learning rate of 0.01, and the discount factor, 7 was set to 0.99.
Table 2: Ablate on Antmaze. The best score metrics is presented in the parenthesis.
Dataset
diff-QL
ours w.o. OPER-R RR
ours
antmaze-umaze-vO
95.6 (96.0)
94.3
94.3 (97.0)
antmaze- mammaze-umaze-ee diverse-v0
(84.0)
88.5
88.5 6.1 (95.0)
antmaze medium- wawddiiimm-ll----- play-v0
0,0 (79.8)
77.0
85.6 1: 1.7 (92.0)
antmaze-mediung mediim-diverse--ddd
6.4 (82.0)
70.4
83.9 1.6 (90.7)
antmaze-large antmaze-large-plaa--u
1.6 (49.0)
47.8
65.4 + 8.6 (74.0)
antmaze- Larazze--lrge-dde lagge-iiveree-
4.4 (61.7)
60.3
67. = 1.8 (75.7)
average
29.6 (75.4)
73.1
80.8 (87.4)
welkerZdd
hopper
5000
Average Retur
4000
3000
2000
2000
SAC EMA
1000
1000
SAC w.o. EMA (LayerNorm)
1000
200
800
1000
steps
sleps
Figure 11: SAC in the online setting.
F More Experiments
F1 Validate the homogeneity in Lemma 1.
We define a family of 3-layer Rel. a-acivvvieeted ML.Ps (with bias term) with different scaling of the same network parameter 0 (from D4RL training checkpoints when value divergence happens), and
22

---

feed these networks with the same imput. We show one example in the below table, which confirms that the L-degree homogeneity of output. gradient and NTK is valid high precision, the NTK is almost parallel, too. The larger the scaling factor is, the more accurate k increasing pattern is for output 1 for gradient and 12(4-1) for NTK). Further, we empirically validated the L-degree homogeneity NTK holds for all checkpoints in D4RL experiments where divergence happens.
Table 3: Validate homogeneity Lemma
10
fao f (x)
124.99 (5
999.78 78 (10")
grad scale
25.001 (5")
100.001 (10")
NTK scale
624.99(5')
9999.84 8999.8400 (10")
NTK cos
1-10-
10
F.2
Ablate the contributions of EMA, Double-Q, and LayerNorm.
We conducted an ablation study in the following sequence: "DDPG without EMA", "DDPG + "DDPG+EMM", "DDPG + EMA + DoubleQ". and "DDPG + EMA + DoubleQ + LayerNorm". Each successive step built upon the last. We carried out these experiments in settings of 10% and 50% Mujoco environments notably susceptible to divergence. We report the total scores across nine tasks in Figure 12. It clearly reveals that the EMA and double-Q, which are used as common pracitce in popular offline RL algorithms such CQL. TD3+BC contribute preventing diverngence and thus boosting performance. Nevertheless, built upon EMA and DoubleQ, our LayerNorm can futher significantly boost the performance by surpressing divergence.
EMA
600
+EMA
*LayerNorm
Performance
400
200
50% Mujoco Dataset
10% Mujoco Dataset
Figure 12: Ablate EMA, Double-Q, and LayerNorm 10% and 50% Mujoco datasets.
F.3 Benchmarking Normalizations.
Previously, we have demonstrated that LayerNorm, BatchNorm, and WeightNorm can effectively maintain low SEEM and stabilize convergence in Section 4. Our next goal is identify the most suitable regularization method for the value network in offline RL. Prior research has shown that divergence correlated with poor control performance[48; 19]. In this context, evaluate the effectiveness of various regularization techniques based on their performance in two distinct settings the Antmaze task and the X% Mujoco dataset we mentioned above. Previous offline RL algorithms have not performed particularly well in these challenging scenarios. As displayed in Figure 13, TD3+BC, when coupled with layer normalization batch normalization. yields significant performance enhancement on the 10% Mujoco datasets. The inability of batch normalization to improve the performance might be attributed to the oscillation issue previously discussed in Section 4. In the case of Antmaze tasks, which contain numerous suboptimal trajectories, we select TD3 with a diffusion policy, namely Diff-QL [51], as our baseline. The diffusion policy capable of capturing multi-modal behavior. As demonstrated in Figure 13 and Table 4, LayerNorm can markedly enhance performance on challenging Antmaze tasks. In summary, we empirically find LayerNorm to be suitable normalization for the critic in offline RL.
23

---

Regelarizations
090
Pafformince
480
400
302
LN
BN
Figure 13: Normalizations effect on 10% Mujoco Locomotion Datasets.
Table 4: Normalizations effect two challenging Antmaze tasks.
Dataset
diff-QI.
I.N
BN
WN
latmaze-large-play-dd antmaze-large-iivrss--o
1.6
72.7
1.0
35.0
4.4
66.5
2.1
42.5
F4 How LayerNorm should be added.
The inclusion of LayerNorm is situated between the linear and activation layers. However, the ideal configuration for adding LayerNorm vary and may depend factors such 1) the specific layers to which LayerNorm should be added, and 2) whether or not to apply learnable per-element affine parameters. To explore these variables, conducted assessment of their impacts performance in the two most challenging Antmaze environments. Our experimental setup mirrored that of the Antmaze experiments mentioned above, utilizing three-layer MIP critic with hidden size configuration of (256,256,250) We evaluated variants where LayerNorm was only applied to a portion of hidden layers and where learnable affine parameters performances with LayerNorm applied solely to the initial layers LN (0), LN (0,1) are considerably lower compared the other setups the geatmaze-argge-laa-voee task, while applying LayerNorm task,
disabled. As in Table the
to all layers LN(0,1,2) seems to yield the best performance. For the "antmaze-large-dveerr performances seem be more consistent across different LayerNorm applications. Overall, this analysis suggests that applying LayerNorm to all layers tends to yield the best performance in these tasks. Also, the utilization of learnable affine parameters appears less critical in this context.
Table The effect of LayerNorm implementations on two challenging Antmaze tasks.
Dataset
w.o. LN
LN
LN
LN
LN
LN
LN (no learnable)
(0)
(0,1.)
(1,2)
(2)
(0,1,2)
Malmazz-large-plaa
1.6
0
0
8.3
17.8
72.7
72.8
antmaze-large-diverr
4.4
60.2
68
77.1
65.5
66.5
66.7
E55 Baird's Counterexample.
Here, we will present our solution allow linear approximation solve Baird's Counterexample which was first used to show the potential divergence of Q-learning. Baird's Counterexample specifically targets linear approximation. Neural networks exhibited stable convergence within the Baird counterexample scenario. Thus, we investigated cases of divergence within linear situations and discovered that incorporating LayerNorm in front linear approximation could effectively lead to convergence (see Figure 14). This finding is not entirely surprising. Since línear regression model can be considered special case of neural networks, our analysis can be directly applied to linear settings.

---

Q-learning on Baird's Counterexample
No LayerNorm LayerNorm
5030
Weight Norm
3030
1030
0
290
500
750
1003
1250
1500
1750
2003
Steps
Figure 14: Baird's Counterexample.
G
Discussion
SEEM and Deadly Triad. Deadly Triad is term that refers to problematic interaction observed in reinforcement learning algorithms, where off-policy learning, function approximation, and boot- strapping converge, leading to divergence during training. Existing studies primarily analyze linear functions as Q-values, which tend to limit the analysis to specific toy examples. In contrast, our work uses NTK theory to provide an in-depth understanding of the divergence of Q-values in non-linear neural networks realistic settings, and introduces SEEM as a tool to depict such divergence. SEEM can be used to understand the Deadly Triad as follows: If a policy is nearly on-policy, X, is merely perturbation of X. Consequently, A, A (X; X) Go, (X. x) z (7 - 1)Go, (X, X). with G' tending to be negative-definite. Without function approximation. the update of Q(X) will not influence Q ( and the first term in A, becomes zero A, = Go, (X, ensures that SEEM non-pon non-positivve and Q-value iteration remains non-expansive. If we avoid bootstrapping, the value iteration transforms into a supervised learning problem with well-understood convergence properties. However, when all three components in Deadly Triad present, the NTK analysis gives rise the form A, = Ge, (Xi, X) Ga, (X,X), which may result in divergence if the SEEM is positive.
Linearity Prediction Outside Dataset Range. As has been pointed out in [54] (Fig. 1), ReLU- activated MLP without any norm layer becomes a linear function for the points that are excessively of the dataset range. This fact can be comprehended in the following way: Consider input Az, the activation state for every neuron in the network becomes deterministic when + co. Therefore, the ReLU activation layer becomes multiplication of constant 0-1 mask, and the whole network degenerates to a linear function. If the network becomes a linear function for inputs A À 2 C, this will cause arbitrarily large NTK value between extreme points Azo and dataset point z0o- Because now So (Xr) can be written in equivalent form W so oi Xu)) AW becomes linear proportional A. This further induces linear NTK value between x and Az that is unbounded. This phenomenon can be intuitively interpreted as below, since the value prediction of fe (z) on line Any is almost linear function W T x for large A, any subtle change of effective parameter W will induce a large change on Azo for large X. This demonstrates a strong correlation between the value prediction between point Zg and far-away extreme point Amy. The rigorous proof can be found in Proposition 1.
Policy Constraint and LayerNorm. We have established a connection between SEEM and value divergence. As shown in Figure 6. policy constraint alone can also control SEEM and prevent divergence. In effect, policy constraint addresses an aspect of the Deadly Triad by managing the degree of off-policy learning. However, an overemphasis on policy constraint, leading to excessive bias, can be detrimental to the policy and impair performance, as depicted in Figure 7. Building on this insight, we focus on an orthogonal perspective in deadly triad regularizing the generalization capacity of the critic network. Specifically, we propose the use of LayerNorm in the critic network to inhibit value divergence and enhance agent performance. Policy constraint introduces an explicit bias into the policy, while LayerNorm does not. Learning useful information often requires some degree of prior bias towards offline dataset, but too much can hinder performance. LayerNorm, offering an orthogonal perspective to policy constraint, aids in striking better balance.
25

---

Action Norm Extreme Norm
antmaae--argge antmaace-bare--llee-e
144
1.0
1.0
1.0
0.6
0.8
0.5
04
02
Policy Action
0.2
0.2
0.2
Data Action
1000
1000
500
00
K. Staps
Shapps
Slipa
Ships
hooppe-medium-22
walke22-mmecuum-
bolfidhetst--meddu--uuao 2
walker2ommedueeepppe--
0.6
0.6
0.6
0.5
250
500
750
1000
280
540
750
1000
0.0
250
500
750 1000
25:
500
750
1000
Stopps
Steps
Figure 15: An observation that policy actions tend to gravitate toward extreme points when Q-value divergence occurs. In D4RL environments, where action spaces are bounded, the ratio the policy action norm the norm of points at the boundary (i.e., extreme points) provides indication of whether the current action is approaching these extreme points. To facilitate the observation of Q-value divergence, the experimental setup employed is consistent with that described in Section of the paper. From the figure, it can be seen that the ratios for policy actions are nearly 1, suggesting close proximity to the extreme points. In contrast, the ratios for actions within the dataset significantly lower than 1.
H More Visualization Results
In Assumption 2, posit that the direction of NTK and the policy remains stable following certain period of training. We validates this assumption through experimental studies. We observe the convergence of the NTK trajectory and policy in all D4RL Mujoco Locomotion and Antmaze tasks, as depicted in the first two columns of Figures Figure 16, Figure 17, and Figure 18. We also illustrate the linear growth characteristic of Adam optimization (as outlined in Theorem Theorem 4) in the fourth column. As a consequence, the model parameter vectors maintain a parallel trajectory, keeping the cosine similarity near shown the third column. Figure 19 and Figure 20 showcase how SEEM serves as a "divergence detector 'in Mujoco and Antamze tasks. The surge in the SEEM value is consistently synchronized with an increase the estimated Q-value.
26

---

welker22-medkmm22 (3 seeds)
NTK simlartty
Action similarity
Model similartty
Model Norm
1.0
1.0
1.0
200000
0.8
0.8
0.5
100000
60000
1000
2000
3000
1000
2000
3020
1000
2000
9000
1000
2000
3000
Staaps
Stops
stopps
Stspp
walker2d-median hetrrr2ddddimm-mppppp/.
seads)
NTK simierity
Action similanly
Model simierily
Modell Norm
160000
100000
0.2
0.2
0.2
0.0
1000
2003
5000
0.0
1000
2000
5000
0.0
1030
2003
5000
1000
2000
3000
<tstpss
stops
stops
stops
walker2d mad 2 similatty
similartty
similatty
200000
15:0000
0.8
0.8
0.5
0.4
0.4
0.4
0.2
0.2
0.2
50000
Kataps
Ktupps
stupe
Figure 16: NTK similarity, action similarity, model parameter similarity, and model parameter norm curves in D4RL Mujoco Walker2d tasks.
27

---

halcheotaa--mduu-222 seeds)
NTK simlirity
similartty
Model similrity
200000
100000
0.8
0.8
0.5
100000
04
0.4
0.4
50000
02
0.2
0.2
0.0
1000
3000
0.0
2000
0.0
9000
2000
Shops
Stops
shapp
haftthadthh-mniiix----laa--y (3 seads)
NTK simllartty
Action similartty
Model similarty
Model Norm
1.0
1.0
1,0
150000
0.8
0.8
0.8
08
0.6
0.8
100000
50003
5000
2000
1030
5000
2000
Ktspps
steps
steps
steps
infccheedd-mmedieerrerriiiz (3 seeds)
NTK simliarity
Action simianty
Model similartty
Model Norm
1.0
1.0
15:0003
100000
60000
0.0
1000
2003
2000
0.0
1000
2000
3000
0.0
1030
2003
3000
1000
2000
3000
stsps
K stops
sstops
K. stops
hopper- um-v2 seeds)
NTK similortty
Action similarity
Model similartty
Modill Norm
200000
15:0003
0.6
0.4
0.4
0.4
0.2
0.2
0.2
0.0
1000
3000
0.0
2000
0.0
1000
8000
2000
ssssps
K steps
steps
K steps
hopper-meddum-eellayy-
seeds)
NTK simitarity
Action similanly
Model simierity
Modell Norm
1.0
1.0
1,0
200000
0.8
0.8
0.8
190000
0.8
0.6
0.6
0.4
0.4
0.4
50000
staps
stops
happor-nacium-con expprttv22 (3 saads)
NTK simliarity
Action simianty
Model similartty
Model Norm
1.0
1.0
1.0
0.8
0.8
0.5
100000
50000
1000
2000
3000
0.8
1000
2000
3000
1030
2000
5000
1000
2000
3000
Callps
steps
steps
stape
Figure 17: NTK similarity, action similarity, model parameter similarity, and model parameter norm curves in D4RL Mujoco Halfcheetah and Hopper tasks.
28

---

v0 seecs)
simlirity
simiarity
similaity
200000
154000
0.8
0.8
0.5
103000
04
0.4
0.4
50000
02
0.2
0.2
0.0
1000
3000
0.0
2000
0.0
2000
<llpps
shaps
Stops
shapp
seads)
NTK simllartty
Action similartty
Model similartty
Model Norm
200000
1.0
1.0
1,0
0.8
0.8
0.8
08
0.6
0.8
100000
52000
5000
2000
1030
2000
Ktspps
steps
steps
steps
animmae-ttedium-llay-- (3 seeds)
NTK simliarity
Action simiantty
Model similartty
Model Norm
1.0
1.0
200003
100000
0.0
0.0
0.0
1000
2003
3000
1000
2000
3000
1030
2003
3000
1000
2000
3000
stspps
stops
stopps
K. stops
medium civerse-vO (3 seeds)
NTK similarty
Action similariy
Model similarit
Model
20000
0.4
0.4
0.4
0.2
0.2
0.2
0.0
1200
2000
3000
0.0
1200
2000
steps
K slaps
K steps
steps
antmeuze- paae
NTK simitrrtty
Action similanly
Model similrrly
Model Norm
1.0
1.0
1.0
0.8
0.8
0.8
100000
0.8
0.6
0.8
75000
0.4
0.4
0.4
60000
stops
antmaze- vntmaze-large--te waree-lang--vviiee pengo-diieesovvvee vo (3 seeds)
NTK similarity
Action similatty
Model similartty
Model Norm
1.0
1.0
1.0
0.8
0.5
103000
75000
25000
1000
2000
3000
1000
2000
3000
1030
2000
5000
1000
2000
3000
Caalps
K staps
steps
stape
Figure 18: NTK similarity, action similarity, model parameter similarity, and model parameter norm curves in Antmaze tasks.
29

---

walker2d mediu X2 (3 seeds)
walker2d medium roplay seeds)
1.0.
SEEM/O vauu
1,0
SEEM'O volub
1.0
SEEMIO volue
SEEMG value
1,0 SEENVO valus
1,0
SEEMO volue
0.5
0.5
10"
0.5
10
0.6
0.5
10
0.5
10
0.0g
2000
0.0
0.0
0.0g
2000
0.0
0.0g
Kaapa
K
Kapp
wwweea--mmiinn.opppr... (3 seads) -wlle
halthhwiah--addaa--222 (3 sunds) SEEM
SEEMIO-vlus
BEE
SEEMC-value
10
SEEMO-vaku
0.5
0.5
0.5
10
0.5
0.5
10
0.5
0.00
2000
0.0g
2000
0.00
2000
0.0g
2000
0.0g
0.0g
2000
X stape
Karap
Kinesa
steps
K. IiXEnce
hofehhetar mww.tan-mmmmaappopommom seeds)
halfchedtahhcommm medimdddn expert-2 (3 seeds]
1.0
SEFMIG-WWI
SEIMIQ--sba
SEEMG-
SEFM
1,0
SFFMQ2-mia
0.5
0.5-
0.5-
0.5
10
0.5-
0.5-
10
0.0
2000
0.0g
2000
0.0g
2000
0.0
0
2000
0.0g
2000
0,0
2000
steaa
zzaaa
Knaaa
stapa
Ktago
hoppar-maiiinm-22
hepper-mediian www.im--immmaaaaoem seeds)
1.0
SE GE
1.0 SE
1.0
1.0
SEEMO-Vilee
1.0
SEEMO-Mi
1.0
SEEM/Q-wkle
10
10
10
0.5
10
0.0
2000
0.0g
2000
0.0
2000
0.0
2000
0.0g
2000
0.0g
2000
wegper-mdduu--pppee-ee 23
1.0
SB EM
SEEM/G--atue
10
0.5
0.5
0.0
2000
0.0g
2000
0.0g
2000
Figure 19: The normalized kernel matrix's 's SEEM (in red) and the estimated Q-value (in blue) in D4RL Mujoco tasks. For each environment, results from three distinct seeds are reported.
30

---

SEE
SCDAG
1E60 2000 1003
9.4g
100
2000
3500
1931
20:0
1516
2010
5000
1030
2030
300
1000
2000
5000
SECMOQ valla
SCTM/G cals
SCOVIO va
SCEMD volua
1.0
1.0
untriic-erreee me00
1.5
CLEAED
1.5
SEEVO
1.0
SESETO
soats
SLEMI
0.5
9.5
0.5
0.
0.5
2002
3003
0.0g
1235
2000
0.5g
1039
2322
300
(a) The normalized kernel matrix's 's SEEM (in red) and the estimated Q-value (in blue) in D4RL Mujoco tasks.
sueds) Unnomalied SEEM
entrraa-rmmaaadiveaea seeds)
Unnorma lzed SEEM
Unnermaizee SEEM
Unnormalzzed SEEM
Unnormalized SEEM
Unnommtzzed SEEM
10
10"
10"
10"
10
10
10
10
2000
2000
2000
2000
2000
2000
Ksuups
wntmore-mediume am-- sends)
antmaze- uidiverne- seeds)
Unnorma lznd SEEM
Undomalzed SEEM
Unnormaiized SEEM
Unnormalzzed zed SEEM
Unnomalizeed 3EEM
Unnormalized SEEM
10
10"
10
10'
10
10'
16"
10
2000
2000
2000
2000
2000
2000 Unnermeized SEEM
antmaze large play v0 seeds)
antmaze large diverte vo 13 seeds)
Unnormatzed SEEM
Undomallzed BEEM
Unnommlizedd BEEM
Unnormalions SEEM
Undommeizeed BEEM
10
10
10
10
10"
10
10
10
10
10
10
AWp
K Kooass
K stupo
Ksluss
K
stups
(b) The unnormalized kernel matrix's SEEM. The three curves in each environment correspond directly to those presented in Figure (a)
Figure 20: In Figure (a), an inflation in the estimated Q-value coincides with surge the normalized SEEM. However, there are some anomalies, such the second running in the 'umaze-diverse' environment, where the Q-value rises while the unnormalized SEEM remains low. However, the corresponding normalized SEEM in Figure (b) suggests actual inflation of SEEM. Furthermore, for scenarios where the Q-value converges, as seen in the second running in 'large-diverse', the unnormalized SEEM maintains an approximate zero value.
31


Reviews:
[{"review_id": "kqRcdE6wY2", "reviewer": "Reviewer_WzD5", "summary": "This paper analyzes Q-value divergence in Offline-RL by considering a\nneural tangent kernel for the value function.  They show that\nconsideration of this kernel is predictive of Q-value divergence.\nThis analysis further leads to the observation that using a LayerNorm\nyields a kernel that behaves more like one would hope -- with nearer\nvalues being more impacted than farther ones.  Their empirical tests\non D4RL benchmarks similarly show the benefits of LayerNorm.\n", "questions": "None", "limitations": "N/A", "rating": 7, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 4, "strengths": "I like the insights in this paper and suspect that its publication\nwill lead others to further engineering that will yield even better\nperforming adjustments than the simple addition of the LayerNorm.\nFigure 5 suggests there is more to be done.  Why does the figure on\nthe left indicate almost the opposite of what one would hope for.  At\na minimum, should the value right at x_0 be red?\n", "weaknesses": "The paper needs a number of editing improvements.  First, it needs an\nEnglish grammar checker.  The errors mostly don't lead to difficulties\nin understanding, but, the paper should not be published with the\ncurrent level of English quality.\n\nIn section 5.1, the BC term should be explained even though the\nreference is given for it.\n\nThe reference to a score of 81.4 in that section does not appear to match\nthe numbers in Table 1.\n\nThe legends in the figures need to use a larger font size.\n\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your time and thoughtful review. On your questions, please see the response below.\n# Explanation for Figure 5\nPlease note that the color in Figure 5 represents the normalized NTK. Indeed, in the left figure, the absolute NTK surrounding $x_0$ is positive. However, values farther from $x_0$ are significantly larger. As a result, the NTK around $x_0$ is minimal and normalized to a value close to zero, as indicated by the blue color.\n# Addressing  Weaknesses\nWe appreciate your suggestions on grammar, font size, and other issues and will correct them in our formal version. We will conduct a thorough grammar check to correct grammar and spelling mistakes. Additionally, we intend to provide a formal introduction to BC and adjust the font size of the legends in subsequent versions of the paper.\n\nThe value 81.4 in Line 289 is a previous result obtained using fewer seeds, whereas 80.8 in the table represents the accurate result computed with 10 seeds. We will rectify this discrepancy. "}, {"review_id": "UYV3V79o36", "reviewer": "Reviewer_T4o6", "summary": "The paper theoretically investigates the problem of value function overestimation in offline RL through the lens of neural tangent kernel (NTK). Additionally, the paper presents empirical findings that validate the effectiveness of incorporating LayerNorm before each activation function in mitigating value network divergence. The paper conducts extensive experiments on D4RL AntMaze and D4RL MuJoCo Gym while varying the dataset size and demonstrates that LayerNorm ensures stable value convergence and leads to state-of-the-art performance even with significantly small datasets.", "questions": "Please see the weaknesses above.", "limitations": "N/A", "rating": 5, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "- The paper is well written and organized. The paper provides empirical evidences to support its theoretical claims (Figures 3, 4).\n- The proposed method can be easily integrates with existing offline RL algorithms (e.g., CQL, IQL, TD3+BC) and consistently improves performance.", "weaknesses": "- It is worth noting that DR3 has already investigated the dynamics of Q-learning in offline RL using NTK [1]. It is crucial for the author to properly reference this prior work and to establish a clear connection between the two studies. The author should explicitly highlight the similarities and differences between the proposed analysis and the findings presented in DR3.\n- The idea of applying LayerNorm (or GroupNorm) has already been proposed in Scaled QL [2].\n- There should be a proper reference to support the claim that the neural network becomes a linear function when the input’s norm is too large (line 234).\n- It is unclear why the linearity of the neural network leads to an improperly large kernel value (line 234).\n- The experimental setup in Section 5.2 is very similar to DR3 [1].\n\n[1] Aviral Kumar et al., DR3: Value-Based Deep Reinforcement Learning Requires Explicit Regularization, ICLR 2022. \\\n[2] Aviral Kumar et al., Offline Q-Learning on Diverse Multi-Task Data Both Scales And Generalizes, ICLR 2023.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your time and thoughtful review. Please see the response below.\n# Linear NTK Value\nPlease refer to the global response for a detailed explanation.\n# About Missing References\nThank you very much for pointing out these two papers that we missed. Since we conducted literature research along the line of works about Q-value divergence and deadly triad in offline RL, we were unaware of these two papers when writing the draft. We apologize for such unintended oversight. \nAfter carefully re-examing these two papers, we agree that these two papers are very relevant to our work, hence we will add them to our related work section with discussions in all future versions. \n\n**However, despite similar observations, we would like to highlight some important differences between our work and these two references.** \n# Comparison with DR3\nIndeed, our work shares similar observations of the connection between Q-value divergence and feature rank collapse (called feature co-adaptation in DR3). However, our work is different in the following aspects.\n## 1. More General Setting & Different Perspectives\nWhile DR3 attributes feature rank collapse to implicit regularization where $L(\\theta)=0$, we propose a different perspective. \nTo start with, we first find that \n* (i) Normalization-free networks' value prediction of the dataset sample and extreme point (i.e. action $A_{ex}$ with maximum value entries) exhibit a strange but strong correlation (having large NTK value) \n* (ii) Such networks also tend to output large values for these $A_{ex}$ (see global response), which makes these extreme points easily become policy actions. In Fig. 3 in global response pdf, you can see policy actions tend to move toward extreme points when Q-value divergence occurs in all D4RL tasks.\n\nIt is then shown that (i) and (ii) are the root cause for many consequences, including feature co-adaptation and Q-value divergence. \n1. The Q-value estimation network overestimates some out-of-sample (typically extreme point) action $A_{ex}$'s Q-value $\\hat{Q}(s', A_{ex})$. \n2. This inflated value becomes bootstrapped target, passing its influence to state $s$ by TD update, where the offline dataset contains transition tuple $(s,a,s')$. The Q-network is updated to boost $\\hat{Q}(s,a)$.\n3. Due to the large NTK value between $(s', A_{ex})$ and $(s,a)$ (see Fig. 5), TD updates the network's parameter to increase $\\hat{Q}(s', A_{ex})$ even more than $\\hat{Q}(s,a)$.\n4. A loop is formed by 2 and 3. $\\hat{Q}(s,a)$ keeps chasing $\\hat{Q}(s', A_{ex})$'s value, leading the model's parameter to diverge along certain direction to infinity, which causes every input's feature becomes parallel. \n\nThis framework is **not restricted to over-parametrized regime which requires $L(\\theta)=0$, nor does it require the label noise $\\varepsilon$ from SGD**. Actually, such near-zero critic loss assumption is mostly not true in the real practice of RL. Also, our analysis does not rely on any n-dimensional basis or sub-space assumptions. \n\n**Therefore, our analysis provides an alternative explanation for feature rank collapse and Q-value divergence from the perspective of norm-free network's pathological extrapolation. It applies to more general settings and provides new information.** These novel insights and implications are not covered in previous work, because it requires several non-trivial observations together with a rigorous theory to encapsulate them. As reviewer 4Xya said, our work \"results in a new viewpoint in handling divergence in Q estimation and could yield many improvements in this domain in the future.\"\n## 2. Divergence Metric & Precise Dynamics\nWe established a measure (SEEM) that can synchronously signify divergence. **SEEM is not just an index to indicate divergence, but also a new handle to solve the problem**. Moreover, DR3 does not accurately characterize model's dynamics as our work does. We formally prove in Theorem 1, 2, 3 that the model's parameter $\\theta$ **evolves linearly for the Adam optimizer** (see Fig 4 and 14), and **collapses at a certain iter for SGD** by solving an ODE (see Fig 3). This accurate prediction demonstrates our precise understanding of neural network dynamics and is absent in previous work like DR3 and Scaled-QL.\n## 3. Simpler and Cheaper\nDR3 requires hyperparameter tuning for $c_0$, which involves more effort in searching proper values across different environments. Moreover, it introduces extra computation when computing the gradient of $\\phi(s',a')$ to get $\\overline{\\mathcal{R}}_{exp}(\\theta)$. Our method is free of all these shortcomings by simple layernorm.\n# Comparison with Scaled-QL\nThank you for mentioning this work. Indeed, as we acknowledged in Line 242-247, some previous works like LB-SAC, RLPD have empirically utilized LayerNorm. We will further add Scaled-QL to reference. Nevertheless, compared to Scaled-QL, our contribution is three-fold:\n## 1. Theoretical Justification\nWe provide rigorous theoretical justification for why normalization helps where previous works do not. **The fundamental mechanism of Q-value divergence and its relation to the extrapolation of Q-networks are studied.** \n## 2. Comprehensive Empirical Evaluation\nThe sole application of Scaled QL on CQL doesn't show LayerNorm's effectiveness for other popular offline RL algorithms, while our research confirms the efficacy of regularization methods that can diminish SEEM and improve performance **across a range of offline RL algorithms like CQL, IQL, TD3+BC, and Diffusion-QL, including the primary classes of offline RL.**\n## 3. Beyond Normalization\nAs the other reviewers acknowledged, the SEEM metric \"will lead others to further engineering that will yield even better-performing adjustments than the simple addition of the LayerNorm\". With the SEEM metrics, **we can examine regularization methods from the DL toolbox and even create new methods.** This could potentially inspire further advancements in the field.\n\n"}, {"review_id": "KKneSvrNJh", "reviewer": "Reviewer_4Xya", "summary": "This paper studies the divergence phenomenon in Q-value iteration methods (e.g. Q-learning), especially focusing on the offline RL scenario. They introduce a theoretical framework for studying this issue, predicting divergence and even the training step at which it is likely to happen. Such analysis, which is based on the NTK (Neural Tangent Kernel), fully contains the linear case study (i.e. deadly triad analysis in linear settings) and extends it to non-linear value approximation. The insights they draw allow them to show that reducing a new metric, SEEM which measures self-excitation using NTK, reduces divergence by regularising the network. Regularising the generalisation of Q-networks as opposed to the dominant view of constraining the policy results in a new viewpoint in handling divergence in Q estimation and could yield many improvements in this domain in the future. They experiment with various regularisation methods from the deep learning toolbox, including dropout, BatchNorm, and LayerNorm, showing that LayerNorm is best able to stabilise training (i.e. avoid blowup in Q estimates) and achieve more homogenous and lower SEEM. They show that BC (behaviour cloning) also achieves a low SEEM but at the cost of bias due to policy constraints, which does not allow it to find the best solution or perform well in more challenging domains.", "questions": "1. Given that this paper does not consider EMA or frozen targets (as in DQN), I am wondering how much using such gradual target updates contributes to alleviating divergence in Q estimation. In relation to the above question, I'm also uncertain how much using double Q-learning alleviates such issues in the offline RL setting. Wouldn't it be useful to have an ablation study to see, e.g., how much adding each of these in combination with each other helps performance and how much each of these help in isolation? Say, use **DDPG w/o EMA** vs. **DDPG w EMA** vs. **DDPG w EMA and Double Q-learning** vs. **DDPG w EMA and Double Q-learning and LayerNorm**, showing that additional advancement improves performance. Or compare in isolation, e.g.: **DDPG w/o EMA w LayerNorm** vs. **DDPG w EMA w/o LayerNorm**.\n\n2. Are the results in Fig. 9 in combination with EMA? If not, what if you use EMA to show how the insights from the scenario w/o EMA carry over to the with-EMA case?  \n\n3. Does using LayerNorm as in your solution have the potential of allowing online DQN-style methods to step away from EMA/frozen target approaches altogether towards using the same online and target networks?\n\n4. Does your solution allow DQN to solve Baird’s Counterexample (which was first used to show the potential divergence of Q-learning with function approximation)? \n\n5. Does the analyses apply also to discrete-action pure critic settings such as DQN?\n\n\n**Minor:**\n\n- Can you elaborate on the way BC introduce explicit bias?\n\n- I assume from the size of the $d_0$ layer that the architecture of the Q-function is action-in, correct? \n\nLine 82: don’t -> do not\n\nLine 119: To summary -> In summary\n\nLine 123: explains -> explain", "limitations": "1. Analysis does not directly apply to the standard practice of using Exponential Moving Average (EMA; as in DDPG) or target-network freezing (as in DQN). \n\n2. It is not clear how much of the contributions would carry over to the online learning setting and what would be the implications of using LayerNorm in online RL.\n\n3. The analysis seems to apply to MLP architectures. It is not directly discussed how much of the contributions would apply to commonly used architectures such as ConvNets, ResNets, etc.", "rating": 7, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "- The regularisation viewpoint on this phenomenon is a new one (as far as I know) yet a natural one. \n\n- The regularisation viewpoint frees us from introducing policy constraints (which introduce bias) and allows us to simply solve the problem using off-the-shelf methods from the deep-learning toolbox.\n\n- The theorems (if proofs are correct, I couldn't verify them personally with certainty) bring highly beneficial theoretical insights to the problem. \n\n- The experiments are remarkable at showing that: (1) SEEM is a powerful metric in signifying divergence, (2) LayerNorm reduces SEEM and more homogeneously does so during training than other regularisation methods, resulting in stable Q estimation, (3) LayerNorm combined with SOTA offline RL methods results in significant improvements, especially under scarce access to offline data.", "weaknesses": "- I believe a short description of offline RL methods used could enhance the exposition of the ideas: policy constraint is not formally introduced because the methods incorporating it are not discussed in any depth.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your time and constructive feedback. Please see the response below.\n# Q1: Clarification on EMA, Double-Q, and LayerNorm's Contributions\nYour suggestion to illustrate the interplay of EMA, Double-Q, and our LayerNorm in mitigating Q-value divergence is particularly insightful. To this end, we conducted an ablation study in the following sequence: \"DDPG without EMA\", \"DDPG + EMA\", \"DDPG + EMA + DoubleQ\", and \"DDPG + EMA + DoubleQ + LayerNorm\". Each successive step built upon the last. We carried out these experiments in settings of 10% and 50% Mujoco - environments notably susceptible to divergence. We report the total scores across nine tasks; please refer to Figure 1(a) in the attached PDF in global response for specific results.\nIt clearly reveals that the EMA and double-Q, which are used as common pracitce in popular offline RL algorithms such as CQL, TD3+BC contribute to preventing diverngence and thus boosting performance. Nevertheless, **built upon EMA and DoubleQ, our LayerNorm can futher significantly boost the performance by surpressing divergence.**\n\n# Q2: Whether results in Fig. 9 use  EMA?\nYou're correct. Fig.9, as well as all experiments in Section 5 (Performance Evaluation), **incorporate EMA**. This is in line with previous practical algorithms such as TD3, CQL, and IQL. We conducted a theoretical analysis without EMA to reduce the complexity of the derivation. Nevertheless, the significant outcomes presented in Section 5 imply that insights derived from scenarios without EMA still apply effectively when EMA is included.\n\n# Q3: LayerNorm Allows Online Methods without EMA\nUpon your suggestion, we tested SAC without EMA in two environments in **online settings**: Hopper and Walker. Please refer to Figure 2 in the attached pdf for the curves of return. Surprisingly, we discovered that **LayerNorm solution allows the SAC without EMA to perform equivalently well as the SAC with EMA**. In contrast, SAC without EMA and LayerNorm behaved aimlessly, maintaining near-zero scores. It reveals LayerNorm (or further regularizations discovered by SEEM later) have the potential of allowing online DQN-style methods to step away from EMA/frozen target approaches altogether towards using the same online and target networks.\n\nYour advice led us to these unexpected yet intriguing findings. These results partially reveals potential implications of SEEM in an online learning setting, as written in the limitations section of your review. We leave SEEM in online setting as future work.\n\n\n# Q4: Baird’s Counterexample\nWe would like to clarify that Baird's Counterexample specifically targets **linear** approximation. In our experiment, we found **neural networks** exhibited stable convergence within the Baird counterexample scenario. Furthermore, we investigated cases of divergence within linear situations and discovered that **incorporating LayerNorm in front of linear approximation could effectively lead to convergence** (see Fig. 1(b) in the attached pdf). This finding is not entirely surprising. Since a linear regression model can be considered a special case of neural networks, our analysis can be directly applied to linear settings. \n\n# Q5: Discrete-action pure Critic Settings\nIndeed, our analyses are not limited to continuous scenarios; they are equally applicable to discrete settings. The theoretical insights can be seamlessly transferred to a discrete setting, requiring only negligible alterations. **We have verified our analysis using a simple discrete-action gridworld with DQN**. Notably, in the gridworld scenario, all conclusions mentioned in the paper are observed, including the Q-value divergence in sync with the rise of SEEM, linear decay of Q-value's inverse with the SGD optimizer, and etc. We've also assessed LayerNorm's effectiveness in curbing divergence in the discrete gridworld. Please do not hesitate to inquire further on this topic.\n\n\n# Response to Minors\n[The way BC introduces explicit bias] **Simply speaking, BC restricts any learned policy to be near the behavior policy. This bias is harmful when the behavior policy is sub-optimal**. An illustration is the example in Fig. 7 In our paper. Specifically, the offline dataset of Antmaze-large-play features behaviors of various qualities. When employing a strong BC (with a large BC coefficient of 3 or 10) to circumvent OOD actions and subsequent divergence, the learned policy is forced to mimic some sub-optimal actions, resulting in poor performance.\n\n[Architecture of the Q-function] You are correct. The Q-function is action-in.\n\n# Conclusion\nIn conclusion, we appreciate the opportunity to address these vital concerns. We hope that our responses sufficiently clarify your question.  We would be grateful for your reconsideration of the **confidence**, in light of the explanations provided. If there are any additional questions about our research, please do not hesitate to reach out."}]

Author rebuttal:
**Note**: We appreciate the constructive feedback from all reviewers. We have prepared a PDF document for reviewers, containing figures about extensive experiments and results that further illustrate and support our response in the rebuttal.

# Why linear function and large NTK value?
Some reviewer raises questions about why ReLU-activated MLP without norm layer becomes a linear function when the input's norm is too large, and why the linearity of the neural network leads to an improperly large kernel value. Here is a more detailed explanation with intuition.
- **Linearity Prediction Outside Dataset Range**

As has been pointed out in [1] (Fig. 1), ReLU-activated MLP without any norm layer becomes a linear function for the points that are excessively out of the dataset range. This fact can be comprehended in the following way: Consider an input $\lambda x$, the activation state for every neuron in the network becomes deterministic when $\lambda \to \infty$. Therefore, the ReLU activation layer becomes a multiplication of a constant 0-1 mask, and the whole network degenerates to a linear function.
- **Large NTK Value between extreme point and dataset sample**

If the network becomes a linear function for inputs $\lambda x_0, \lambda\geq C$, this will cause an arbitrarily large NTK value between extreme points $\lambda x_0$ and dataset point $x_0$. Because now $f_{\theta}(\lambda x)$ can be written in equivalent form $W^T (\lambda x)$, so $\phi(\lambda x_0)=\nabla_{W}W^T (\lambda x_0)= \lambda x_0$ becomes linear proportional to $\lambda$. This further induces linear NTK value between $x$ and $\lambda x$, which can be unbounded.

This phenomenon can be intuitively interpreted as below, since the value prediction of $f_{\theta}(x)$ on line $\lambda x_0$ is almost a linear function $W^T x$ for large $\lambda$, any subtle change of effective parameter $W$ will induce a large change on $\lambda x_0$ for large $\lambda$, proportional to $||x||_2$. This demonstrates a strong correlation between the value prediction between point $x_0$ and far-away extreme point $\lambda x_0$. The rigorous proof can be found in Proposition 1.

Note that such strong correlation is not restricted to rigorous paralleled inputs $x_0$ and $\lambda x_0$, but also between many dataset small-norm input $x_0$ and extreme large-norm input $x_{ex}$, once they share similarly activated neurons and effective linear parameters. The so-called "extreme" points' norm does not need to be several magnitudes larger than dataset point $x_0$ either. It just needs to be constantly larger to be able to start the self-excitation procedure, therefore starting the divergence loop.


[1] How Neural Networks Extrapolate: From Feedforward To Graph Neural Networks. Xu et al.



Meta-review: This paper brings an insightful analysis on the Q-value estimation divergence in offline RL, proposes a new metric based on NTK, SEEM, to measure the divergence, and suggests using regularization to mitigate the problem, among which LayerNorm shows best performance in the experiments.

All reviewers appreciate the insights from their theoretical analysis. Empirical evaluation shows the efficacy of the new metric and justifies the benefit of LayerNorm in Q-value estimation.

The main concerns come from the following aspects:
- Review T4o6 points out several flaws in the proof of theory and lemma. After discussion, all agree that the problem can be fixed by clarifying additional assumptions and improving the rigorousness of the proof, which authors have promised to add in the revision.
- Similarity to the work of DR3 (Aviral Kumar et al., ICLR 2022) both of which conduct analysis on the dynamics of Q-value estimation in offline RL and suggest explicit regularization as a solution. The authors first admit the ignorance of that reference in the initial submission and then explain the different perspective of the analysis in the two works and different assumptions applied. Review T4o6 agree with the comparison though remain suspicious if the assumptions used in DR3 is more restrictive than this work. It would be very important for the authors to provide detailed comparison in the revision and, if possible, include DR3 as an additional baseline in the experiments.
- LayerNorm has been used in previous works such as Scaled-QL (Aviral Kumar et al., ICLR 2023) and those already mentioned in the submission. The authors explain that this submission provides theoretical justification to the use of LayerNorm besides other regularizations in contrast to empirical usecases in other works and they provide more comprehensive empirical evaluation.

After the discussion, all reviewers agree on acceptance of this submission. The remaining concerns should not impact the theoretical contribution of this work after the additional promised edits. I strongly encourage the authors to take into consideration all the discussion in the reviewing process, fix the proofs and make a final revision according accordingly.
Decision: Accept (poster)