Title: Understanding, Predicting and Better Resolving Q-Value Divergence in Offline-RL
Venue: NeurIPS 2023
Accepted: yes
Keywords: Offline RL

Theory
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-Value Divergence in Offline-RL

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 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, self-excitation , as the primary cause of Q-value estimation divergence in offline RL. Then, we propose a novel S elfE xcite E igenvalue M easure (SEEM) metric based on Neural Tangent Kernel (NTK) to measure the evolving property of Q-network at training, which provides an intriguing explanation on 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 regularizing the neural network's generalization 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. Code can be found at https://offrl-seem.github.io .

1 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-network [19]) are proposed to 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 a result, directly employing off-policy algorithms in offline settings confronts substantial issues related to value divergence [14; 30].

This raises a 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

identified to be the main cause of value divergence. Specifically, it points out that a RL algorithm is susceptible to divergence at training when three components are combined, including off-policy learning, function approximation, and bootstrapping. As a special case of off-policy RL, offline RL further attributes the divergence to distributional shift : when training a Q-value function using the Bellman operator, the bootstrapping operation frequently queries the value of actions that are unseen in the dataset, resulting in cumulative extrapolation errors [14; 30; 33]. 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-policy degree and bootstrapping-two components of the deadly triad-they often overlook the aspect of function approximation. This neglect leaves several questions unanswered and certain limitations unaddressed in current methodologies.

Howdoes 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 a more challenging scenario is encountered, e.g. , the transitions are very scarce. Instead, we are interested in solutions without detrimental bias.

Our study offers a novel perspective on the challenge of divergence in Offline-RL. We dissect the divergence phenomenon and identify self-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 to 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 is the first work, to our knowledge, that accurately characterizes neural network dynamics in offline RL and explains the divergence phenomenon. More specifically, our contributions are three-fold:

Explanation: We offer a detailed explanation of Q-value estimation divergence in offline RL settings with neural networks as non-linear estimators. We propose a novel metric, the S elfE xcite E igenvalue M easure (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.

Prediction: We can foresee the divergence through SEEM at 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 is 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, we suggest mitigating divergence by regularizing the model'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. Then we conduct experiments to 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 ⊆ R d S is the state space and A ⊆ R d A is the action space. P : S × A ↦→ ∆( S ) is the transition dynamics mapping from state-action pair to distribution in state space. r : S × A ↦→ R is the reward function and ν ∈ ∆( S ) is the initial state distribution. An offline RL algorithm uses an offline dataset

0

250

500

750

1000

K steps

0.0

0.2

0.4

0.6

0.8

1.0

NTK similarity

0

250

500

750

1000

K steps

0.0

0.2

0.4

0.6

0.8

1.0

Action similarity

walker2d-medium-v2 (3 seeds)

0

250

500

750

1000

K steps

0.0

0.2

0.4

0.6

0.8

1.0

NTK similarity

0

250

500

750

1000

K steps

0.0

0.2

0.4

0.6

0.8

1.0

Action similarity

D = { ( s i , a i , s ′ i , r i ) } M i =1 , which consists of M finite samples from interactions with MDP M . We write s ∈ D if there exists some k ∈ [ M ] such that s k = s , similar for notation ( s, a ) ∈ D .

In this paper, we consider using a L -layer ReLU-activated MLP as the Q-value approximator. Each layer has parameter weight W ℓ ∈ R d ℓ +1 × d ℓ and bias b ℓ ∈ R d ℓ +1 . The activation function is ReLU: σ ( x ) = max( x, 0) . Integer d ℓ represents the dimensionality of the ℓ th hidden layer, where d 0 = d S + d A and d L = 1 since a scalar estimation is required for Q-value. Denote θ = Vec[ W 0 , b 0 , . . . , w L , b L ] to be the vector of all learnable parameters. We use ˆ Q θ ( s, a ) = f θ ( s, a ) to denote the estimated Q-value using neural network f θ with parameter θ . Also, ˆ π θ ( s ) = arg max a ˆ Q θ ( s, a ) is its induced policy. Sometimes we will write f θ ( X ) ∈ R M where X = [( s 1 , a 1 ) , . . . , ( s M , a M )] stands for the concatenation of all inputs ( s i , a i ) ∈ D in dataset.

Considering a Q-learning-based offline RL algorithm that alternates between the following two steps:

where η 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, we 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 as a special class of kernel function, which measures the similarity between x and x ′ , with dependency on parameter θ . Given a neural network f θ and two points x , x ′ in the input space, the NTK is defined as k NTK ( x , x ′ ) := ⟨∇ θ f θ ( x ) , ∇ θ f θ ( x ′ ) ⟩ , where the feature extraction is ϕ θ ( x ) := ∇ θ f θ ( x ) . If we have two batches of input X , X ′ , we can compute the Gram matrix as G θ ( X , X ′ ) = ϕ θ ( X ) ⊤ ϕ θ ( X ′ ) .

Intuitively speaking, NTK measures the similarity between two inputs through the lens of a non-linear network f θ . Such similarity can be simply understood as the correlation between the network's value predictions of these two inputs. If k NTK ( 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 θ , say, from θ to θ ′ = θ +∆ θ , denote the prediction change for x as ∆ f ( x ) = f θ ′ ( x ) -f θ ( x ) and for x ′ as ∆ f ( x ′ ) = f θ ′ ( x ′ ) -f θ ( x ′ ) . ∆ f ( x ) and ∆ f ( x ′ ) will be highly correlated, and the strength of the correlation is proportional to k NTK ( x , x ′ ) . This intuition can be explained from simple Taylor expansion, which gives ∆ f ( x ) = ∆ θ ⊤ ϕ θ ( x ) and ∆ f ( x ′ ) = ∆ θ ⊤ ϕ θ ( x ′ ) . These two quantities will be highly correlated if ϕ θ ( x ) and ϕ θ ( x ′ ) are well-aligned, which corresponds to large NTK value between x and x ′ . To summary, k NTK ( x , x ′ ) characterizes the subtle connection between f θ ( x ) and f θ ( x ′ ) caused by neural network'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 a single step of Q-value iteration in Equation (1). Our result is the following theorem.

Theorem 1. Suppose that the network's parameter at iteration t is θ t . For each transition ( s i , a i , s i +1 , r i ) in dataset, denote r = [ r 1 , . . . , r M ] ⊤ ∈ R M , ˆ π θ t ( s ) = arg max a ˆ Q θ t ( s, a ) . Denote x ∗ i,t = ( s i +1 , ˆ π θ t ( s i +1 )) . Concatenate all x ∗ i,t to be X ∗ t . Denote u t = f θ t ( X ) -( r + γ · f θ t ( X ∗ t )) to be TD error vector at iteration t . The learning rate η is infinitesimal. When the maximal point of ˆ Q θ t is stable as t increases, we have the evolving equation for u t +1 as u t +1 = ( I + η A t ) u t , where A t = ( γϕ θ t ( X ∗ t ) -ϕ θ t ( X )) ⊤ ϕ θ t ( X ) = γ G θ t ( X ∗ t , X ) -G θ t ( X , X ) .

G is defined in Section 2. Detailed proof is left in Appendix B. Theorem 1 states that when the policy action ˆ π θ t ( s ) that maximizes Q-values keep stable, leading to invariance of X ∗ t , each Q-value iteration essentially updates the TD error vector by a matrix A t , which is determined by a discount factor γ and the NTK between X and X ∗ t .

3.1 Linear Iteration Dynamics and SEEM

Although Theorem 1 finds that TD update is a linear iteration under certain conditions, the dynamic is still complex in general. The policy ˆ π θ t ( s ) that maximizes Q-values may constantly fluctuate with θ t over the course of training, leading to variations of X ∗ t . Also, the kernel G θ ( · , · ) has dependency on parameter θ . This causes non-linear dynamics of u t and the model parameter vector θ . However, according to our empirical observation, we discover the following interesting phenomenon.

Assumption 2. (Existence of Critical Point) There exists a time step t 0 , a terminate kernel ¯ k ( · , · ) and stabilized policy state-action ¯ X ∗ such that ∥ ∥ ∥ k t ( x , x ′ ) ∥ x ∥·∥ x ′ ∥ -¯ k ( x , x ′ ) ∥ x ∥·∥ x ′ ∥ ∥ ∥ ∥ = o (1) and ∥ ∥ X ∗ t +1 -X ∗ t ∥ ∥ = o ( η ) for any x , x ′ ∈ S × A when t > t 0 . We also assume that ∥ ∥ X ∗ t -¯ X ∗ ∥ ∥ = o (1) .

This assumption states that the training dynamics will reach a steady state, where the policy ˆ π θ t stabilizes and NTK converges to a terminate direction. Usually, constant NTK requires the width of the network to be infinite[21]. But in our experiments 1 , the convergence of NTK is still observed in all our environments despite finite width (see Figure 1). Results of more environments are attached in the 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 ˆ π θ t converges after a certain timestep t 0 . Moreover, we find that ˆ π θ t 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 a steady state.

0

200

400

600

Steps

0.0

0.2

0.4

0.6

0.8

1/Q

y = -0.00144 * x + 0.87

R-squared = 1.00000

1/Q v.s. step

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.

Theorem 3. Suppose we use SGD optimizer for Q-value iteration with learning rate η to be infinitesimal. Given iteration t > t 0 , and A = γ ¯ G ( ¯ X ∗ , X ) -¯ G ( X , X ) , where ¯ G is the Gram matrix under terminal kernel ¯ k . The divergence of u t is equivalent to whether there exists an eigenvalue λ of A such that Re( λ ) > 0 . If converge, we have u t = ( I + η A ) t -t 0 · u t 0 . Otherwise, u t becomes parallel to the eigenvector of the largest eigenvalue λ of A , and its norm diverges to infinity at following order

for some constant C ′ to be determined and L is the number of layers of MLP. Specially, when L = 2 , it reduces to O ( 1 1 -C ′ ληt ) .

Theorem 3 predicts the divergence by whether λ max ( A ) is greater to 0. We term this divergence detector λ max ( A ) as S elfE xcite E igenvalue M easure, or SEEM. Essentially, Theorem 3 states that

0.99

0.95

0.93

0.929 0.928 0.927 0.926 0.925

0.92

0.0

discount factor

0.0

0.1

0.2

0.3

SEEM

walker2d-medium-expert-v2

Q diverge

Q converge

0

500

1000

K steps

0.0

0.2

0.4

0.6

0.8

SEEM

0

500

1000

K steps

0.0

0.2

0.4

0.6

10

1

10

4

10

7

10

10

10

2

10

6

10

10

Q-value

we can monitor SEEM value to know whether the training will diverge. To validate this, we conduct experiments with different discount factor γ . The first term in A is linear to γ , thus larger γ is more susceptible to diverging. This is exactly what we find in reality. In Figure 3, we can see that experiments with large γ 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 a terminating timestep T = 1 C ′ λη . 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 it is true. The terminal Q-value prediction'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 a linear regression, obtaining the green dotted line 1 / ∥ u t ∥ ≈ -1 . 44 × 10 -3 t +0 . 87 , which predicts a singular point at 605 th 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 a 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 t > t 0 , the model will diverge if and only if λ max ( A ) > 0 . If it diverges, we have ∥ θ t ∥ = η √ Pt + o ( t ) and ∥ u t ∥ = Θ( t L ) where P and L are the number of parameters and the number of layers for network f θ , respectively.

Again, the detailed proof is left in the supplementary. Theorem 4 indicates that with a Adam optimizer, the norm of the network parameters grows linearly and the predicted Q-value grows as a polynomial of degree L along the time after a critical point t 0 . We verify this theorem in D4RL environments in Figure 4. We can see that the growth of the norm ∥ θ t ∥ exhibits a straight line after a critic point t 0 . Moreover, we can see log Q ≈ 3 log 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 RL.

3.3 Discussions

Interpretation of SEEM and explain the self-excitation The intuitive interpretation for λ max ( A ) is as below: Think about only one sample x = ( s 0 , a 0 ) and its stabilized next step state-action

0

250

500

750

1000

K steps

0

20000

40000

60000

80000

The norm of Q network

5.00

5.25

5.50

5.75

6.00

log(steps)

9

10

11

12

13

y = 2.88 * x -4.36

R-squared = 0.99449

log(Q) v.s. log(t)

walker2d-medium-v2 (3 seeds)

0

250

500

750

1000

K steps

0

20000

40000

60000

80000

The norm of Q network

5.00

5.25

5.50

5.75

6.00

log(steps)

9

10

11

12

13

y = 3.05 * x -5.48

R-squared = 0.99778

log(Q) v.s. log(t)

x ∗ = ( s 1 , a 1 ) . Without loss of generality, we assume f θ t ( x ) < r + γf θ t +1 ( x ∗ ) . If SEEM is positive, we have γ G θ t ( x ∗ , x ) is larger than G θ t ( x , x ) . Recall that G θ t ( x ∗ , x ) depicts the strength of the bond between x and x ∗ because of generalization. So we know that, whren updating the value of f θ ( x ) towards r + γf θ t +1 ( x ∗ ) , the Q-value iteration inadvertently makes f θ t +1 ( x ∗ ) increase even more than the increment of f θ t +1 ( x ) . Consequently, the TD error r + γf θ ( x ∗ ) -f θ ( x ) expands instead of reducing, due to the target value moving away faster than predicted value, which encourages the above procedure to repeat. This forms a positive feedback loop and causes self-excitation. Such mirage-like property causes the model's parameter and its prediction value to diverge.

Stability of policy action ˆ π θ t ( s ) . Previous Theorem 1 and Theorem 3 rely on the stability of the policy action ˆ π θ t ( s ) after t 0 , and here we further elucidate why such stability occurs. Critic networks without normalization have a tendency to output large values for extreme points, i . e ., action A ex at the boundary of action space (see detailed explanation in 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 ˆ π θ t ( s ) . Consequently, as aforementioned, a loop is formed: ˆ Q ( s, a ) keeps chasing ˆ Q ( s ′ , A ex ) , 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 a special case for neural networks, our analysis is also directly applicable to linear settings by plugging in G ( X , 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 in 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 connection 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 , 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 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 θ evolves linearly for the Adam optimizer, and collapses at a 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 c 0 and also introduces extra computation when computing the gradient of ϕ ( s ′ , a ′ ) to get R exp ( θ ) . 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-distribution data points. To gain an intuitive grasp, we visualize the NTK value in a simple 2-dimensional input space for a two-layer MLP in Figure 5. We designate a reference point x 0 = (0 . 1 , 0 . 2) and calculate the NTK value G θ ( x 0 , x ) for a range of x in the input domain. The heatmap displays high values at the boundaries of the input range,

normalizedNTKheatmapw/oLN

1.0

1.0

4

4

0.8

0.8

2

2

0.6

0.6

0

0

0.4

0.4

2

0.2

0.2

4

-4

2

0

2

4

0.0

-2

0

2

0.0

x1

x1

which suggests that even a slight increase (or decrease) in the prediction value at x 0 will result in an amplified increase ( e.g. , 4.5 × ) in the prediction value at a distant point( e.g. , (4 , 4) ). Note that indeed 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. Thus, the value 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.

This abnormal behavior contradicts the typical understanding of a '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 to avoid divergence: regularizing the model's generalization on out-of-distribution predictions . The main reason for such 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, a simple method to accomplish this would be to insert a LayerNorm prior to each nonlinear activation. We conduct a similar visualization by equipping the MLP with LayerNorm [3]. As shown in Figure 5 right, the value reaches its peak at x 0 and diminishes as the distance grows, demonstrating excellent local properties for well-structured kernels.

Such architectural change is beneficial for controlling the eigenvalue of A by reducing the value of improperly large entries since matrix inequality tells us λ max ( A ) ≤ ∥ A ∥ F . Therefore, it is expected that this method can yield a 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 a 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: 1) 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 in Section Section 5. walker2d-medium-expert-v2

0

250

500

750

1000

K steps

10

4

10

7

10

10

10

13

Estimated Q-value

0

250

500

750

1000

K steps

0.0

0.2

0.4

0.6

0.8

1.0

SEEM

0

250

500

750

1000

K steps

10

2

10

3

Estimated Q-value

0

250

500

750

1000

K steps

0.00

0.01

0.02

0.03

0.04

SEEM

TD3 w.o. EMA

Weight Norm

Weight Decay

Dropout

Behavior Cloning

LayerNorm

BatchNorm

To validate the effectiveness of LayerNorm in practical offline RL settings, we select the walker2dmedium-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

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. WeightNorm 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 a 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].

5 Agent Performance

Previous state-of-the-art 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 on 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 or 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 to capture multi-modal behavior [51], continues to struggle with instability, leading to inferior performance.

We first show policy constraint (BC) is unable to 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 to diverge, and performance drops to 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 of 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,

0

200

400

600

800

1000

K steps

0

2

4

6

8

10

Estimated Q-value

BC coeficient

0.5

1

3

10

0.5 (ours)

0

200

400

600

800

1000

K steps

0.0

0.2

0.4

0.6

0.8

1.0

Performance of Antmaze-large-play

BC coeficient

0.5

1

3

10

0.5 (ours)

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 a thorough hyperparameter search for the coefficient, it was observed that it had no substantive impact on enhancing the stability and performance of diff-QL.

5.2 X % Offline Dataset

Past popular offline RL algorithms primarily require transitions rather than full trajectories for training [40; 31; 26; 12; 7; 51]. 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 s t +1 will be directly updated when sampling { s t +1 , a t +1 , s t +2 , r t +1 } . On the one hand, working with offline datasets consisting solely of transitions enables a more authentic evaluation of an algorithm's stability and generalization capabilities. On the other hand, it is 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 a transition-based scenario. We construct transitionbased datasets by randomly sampling varying proportions ( X %) from the D4RL Mujoco Locomotion datasets. Here, we set several levels for X ∈ { 1 , 10 , 50 , 100 } . When X equals 100, it is equivalent to training on the original D4RL dataset. As X decreases, the risk of encountering out-of-distribution states and actions progressively escalates. To ensure a 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 comparison, we report the total score on nine Mujoco locomotion tasks and the standard deviation over 10 seeds. Figure 8 reveal a 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

100%

50%

10%

1%

0

100

200

300

400

500

600

700

800

Toal Score

TD3+BC

IQL

diff-QL

CQL

about only 50 to 150 total points on nine tasks. Further, we observed value divergence in all algorithms, even for algorithms based on policy constraints or conservative value estimation.

Subsequently, we demonstrate the effectiveness of LayerNorm in improving the poor performance in X%datasets. By adding LayerNorm to the critic, value iteration for these algorithms becomes nonexpansive, ultimately leading to stable convergence. As Figure 9 depicts, under 10% and 1% dataset, all baselines are greatly improved by LayerNorm. For instance, under the 1% dataset, LayerNorm

enhances the performance of CQL from 50 points to 300 points, marking a 6x improvement. Similarly, TD3+BC improves from 70 to 228, representing approximately a 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 is shown to be more effective than relying solely on policy constraints. Another notable observation from Figure 9 is that with sufficient data (the 100% dataset) where previous SOTA algorithms 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 or plentiful.

100%

50%

10%

1%

25

0

50

100

150

200

250

TD3+BC

100%

50%

10%

1%

IQL

100%

50%

10%

1%

diff-QL

100%

50%

10%

1%

CQL

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 or theoretically shown the effect of EMA in stabilizing Q-value and prevent divergence in online setting. To answer the question, We tested whether LayerNorm can replace EMA to prevent value divergence. We conducted experiments in 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. 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. 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 of offline RL and gave a comprehensive 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 on 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 our 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 not 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.

References

Joshua Achiam, Ethan Knight, and Pieter Abbeel. Towards characterizing divergence in deep q-learning. arXiv preprint arXiv:1903.08894 , 2019. 14

Gaon An, Seungyong Moon, Jang-Hyun Kim, and Hyun Oh Song. Uncertainty-based offline reinforcement learning with diversified q-ensemble. NIPS , 2021. 2, 8

Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E Hinton. Layer normalization. arXiv preprint arXiv:1607.06450 , 2016. 7

Leemon Baird. Residual algorithms: Reinforcement learning with function approximation. In Machine Learning Proceedings . 1995. 1, 2, 14

Philip J Ball, Laura Smith, Ilya Kostrikov, and Sergey Levine. Efficient online reinforcement learning with offline data. In ICML , 2023. 7, 14

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

David Brandfonbrener, Will Whitney, Rajesh Ranganath, and Joan Bruna. Offline rl without off-policy evaluation. NIPS , 2021. 9, 14

Huayu Chen, Cheng Lu, Chengyang Ying, Hang Su, and Jun Zhu. Offline reinforcement learning via high-fidelity generative behavior modeling. In ICLR , 2023. 2, 8, 14

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

Justin Fu, Aviral Kumar, Ofir Nachum, George Tucker, and Sergey Levine. D4rl: Datasets for deep data-driven reinforcement learning, 2020. 4, 9

Justin Fu, Aviral Kumar, Matthew Soh, and Sergey Levine. Diagnosing bottlenecks in deep q-learning algorithms. In ICML , 2019. 14

Scott Fujimoto and Shixiang Shane Gu. A minimalist approach to offline reinforcement learning. NIPS , 2021. 2, 8, 9, 14

Scott Fujimoto, Herke Hoof, and David Meger. Addressing function approximation error in actor-critic methods. In ICML , 2018. 1, 10, 14

Scott Fujimoto, David Meger, and Doina Precup. Off-policy deep reinforcement learning without exploration. In ICML , 2019. 1, 2, 14

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. 2, 8

Caglar Gulcehre, Ziyu Wang, Alexander Novikov, Thomas Paine, Sergio G´ omez, Konrad Zolna, Rishabh Agarwal, Josh S Merel, Daniel J Mankowitz, Cosmin Paduraru, et al. Rl unplugged: A suite of benchmarks for offline reinforcement learning. NIPS , 2020. 9

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

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

Hado Hasselt. Double q-learning. NIPS , 2010. 1, 14, 23

Sergey Ioffe and Christian Szegedy. Batch normalization: Accelerating deep network training by reducing internal covariate shift. In ICML , 2015. 7, 14

Arthur Jacot, Franck Gabriel, and Cl´ ement Hongler. Neural tangent kernel: Convergence and generalization in neural networks. NIPS , 2018. 3, 4

Michael Janner, Qiyang Li, and Sergey Levine. Offline reinforcement learning as one big sequence modeling problem. In NIPS , 2021. 14

Natasha Jaques, Asma Ghandeharioun, Judy Hanwen Shen, Craig Ferguson, ` Agata 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

Bingyi Kang, Xiao Ma, Chao Du, Tianyu Pang, and Shuicheng Yan. Efficient diffusion policies for offline reinforcement learning. In NIPS , 2023. 22

Bingyi Kang, Xiao Ma, Yirui Wang, Yang Yue, and Shuicheng Yan. Improving and benchmarking offline reinforcement learning algorithms, 2023. 7, 14, 22

Ilya Kostrikov, Ashvin Nair, and Sergey Levine. Offline reinforcement learning with implicit q-learning. In NIPS Deep RL Workshop , 2021. 8, 9, 14

Alex Krizhevsky, Ilya Sutskever, and Geoffrey E Hinton. Imagenet classification with deep convolutional neural networks. NIPS , 2012. 15

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

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

Aviral Kumar, Justin Fu, Matthew Soh, George Tucker, and Sergey Levine. Stabilizing off-policy q-learning via bootstrapping error reduction. NIPS , 2019. 1, 2, 14

Aviral Kumar, Aurick Zhou, George Tucker, and Sergey Levine. Conservative q-learning for offline reinforcement learning. In NIPS , 2020. 2, 9, 14

Sascha Lange, Thomas Gabel, and Martin Riedmiller. Batch reinforcement learning. In Reinforcement learning . Springer, 2012. 1, 14

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. 2

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

Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. In ICLR , 2019. 7

Rui Lu, Gao Huang, and Simon S. Du. On the power of multitask representation learning in linear mdp, 2021. 14

Rui Lu, Andrew Zhao, Simon S. Du, and Gao Huang. Provable general function class representation learning in multitask bandits and mdps. In NIPS , 2022. 14

Xiao Ma, Bingyi Kang, Zhongwen Xu, Min Lin, and Shuicheng Yan. Mutual information regularized offline reinforcement learning. In NIPS , 2023. 2, 14

Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Andrei A Rusu, Joel Veness, Marc G Bellemare, Alex Graves, Martin Riedmiller, Andreas K Fidjeland, Georg Ostrovski, et al. Human-level control through deep reinforcement learning. nature , 2015. 1, 14

Ashvin Nair, Abhishek Gupta, Murtaza Dalal, and Sergey Levine. Awac: Accelerating online reinforcement learning with offline datasets. arXiv preprint arXiv:2006.09359 , 2020. 9, 14

Alexander Nikulin, Vladislav Kurenkov, Denis Tarasov, Dmitry Akimov, and Sergey Kolesnikov. Q-ensemble for offline rl: Don't scale the ensemble, scale the batch size. arXiv preprint arXiv:2211.11092 , 2022. 2, 7, 8, 14

Xue Bin Peng, Aviral Kumar, Grace Zhang, and Sergey Levine. Advantage-weighted regression: Simple and scalable off-policy reinforcement learning. arXiv preprint arXiv:1910.00177 , 2019. 2, 14

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

Tim Salimans and Durk P Kingma. Weight normalization: A simple reparameterization to accelerate training of deep neural networks. NIPS , 2016. 7

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

Richard S Sutton and Andrew G Barto. Reinforcement learning: An introduction . MIT press, 2018. 1, 14

JN Tsitsiklis and B Van Roy. An analysis of temporal-difference learning with function approximationtechnical. Rep. LIDS-P-2322). Lab. Inf. Decis. Syst. Massachusetts Inst. Technol. Tech. Rep , 1996. 1, 2, 14, 17

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.02648 , 2018. 1, 14, 23

Chaofei Wang, Jiayu Xiao, Yizeng Han, Qisen Yang, Shiji Song, and Gao Huang. Towards learning spatially discriminative feature representations. In ICCV , 2021. 15

Chaofei Wang, Qisen Yang, Rui Huang, Shiji Song, and Gao Huang. Efficient knowledge distillation from model checkpoints. NIPS , 2022. 15

Zhendong Wang, Jonathan J Hunt, and Mingyuan Zhou. Diffusion policies as an expressive policy class for offline reinforcement learning. In ICLR , 2023. 2, 8, 9, 14, 23

Ziyu Wang, Alexander Novikov, Konrad Zolna, Josh Merel, Jost Tobias Springenberg, Scott E. Reed, Bobak Shahriari, Noah Y. Siegel, C ¸ aglar G¨ ulc ¸ehre, Nicolas Heess, and Nando de Freitas. Critic regularized regression. In NIPS , 2020. 14

Yifan Wu, George Tucker, and Ofir Nachum. Behavior regularized offline reinforcement learning. arXiv preprint arXiv:1911.11361 , 2019. 2, 14

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

Qisen Yang, Shenzhi Wang, Matthieu Gaetan Lin, Shiji Song, and Gao Huang. Boosting offline reinforcement learning with action preference query. In ICML , 2023. 14

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.01448 , 2023. 14

Tianhe Yu, Aviral Kumar, Rafael Rafailov, Aravind Rajeswaran, Sergey Levine, and Chelsea Finn. Combo: Conservative offline model-based policy optimization. NIPS , 2021. 2

Yang Yue, Bingyi Kang, Xiao Ma, Gao Huang, Shiji Song, and Shuicheng Yan. Offline prioritized experience replay. arXiv preprint arXiv:2306.05412 , 2023. 22

Yang Yue, Bingyi Kang, Xiao Ma, Zhongwen Xu, Gao Huang, and Shuicheng Yan. Boosting offline reinforcement learning via data rebalancing. In Offline RL Workshop , 2022. 22

Yang Yue, Bingyi Kang, Zhongwen Xu, Gao Huang, and Shuicheng Yan. Value-consistent representation learning for data-efficient reinforcement learning. In AAAI , 2023. 14

Shangtong Zhang, Hengshuai Yao, and Shimon Whiteson. Breaking the deadly triad with a target network. In ICML , 2021. 10

A Related Works

Value Divergence with Neural Network. In online reinforcement learning (RL), off-policy algorithms that employ value function approximation and bootstrapping can experience value divergence, a phenomenon known as 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; 13]. Achiam et al . [1] analyze a 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; 30]. Much of the focus in the field of offline RL has been on controlling the extent of off-policy 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 a 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 contribution extends to explaining divergence and proposing promising solutions from the perspective of reducing SEEM. Specially, we 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 on 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-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-QL [51] further replace the BC loss with a diffusion loss and using diffusion models as the policy. AWR-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-distribution 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 f θ to λ times, namely θ ′ = λ θ , then we have following equations hold for any x , x ′ ∈ S × A

Proof . Recursively define

Then it is easy to see that if we multiply each W ℓ and b ℓ by λ , denote the new corresponding value to be z ′ ℓ and note that the bias term is negligible, we have

Note that λ ∥ w ℓ δ ℓ ∥ = λ ∥ W ℓ ( σ ( λ ℓ z ℓ + ϵ ℓ ) -σ ( z ′ ℓ +1 )) ∥ ≤ λ ∥ W ℓ ∥ · ∥ ( λ ℓ z ℓ + ϵ ℓ ) -z ′ ℓ +1 ∥ , the last step is because σ ( · ) is 1-Lipshitz, and we have ( λ ℓ z ℓ + ϵ ℓ ) -z ℓ +1 = ϵ ℓ and ∥ ϵ ℓ ∥ = o ( λ ℓ ) . When value divergence is about to occur, the dot product W i z i becomes relatively large compared to the scalar b i , and the effect of bias on the output becomes negligible. Therefore, we can recursively deduce that the remaining error term

is always negiligle compared with z ′ ℓ +1 . Hence we know f θ ′ ( x ) is close to λ L f θ ( x ) with only o (1) relative error.

Taking gradient backwards, we know that ∥ ∥ ∥ ∂f ∂ W ℓ ∥ ∥ ∥ is proportional to both ∥ ˜ z ℓ ∥ and ∥ ∥ ∥ ∂f ∂ z ℓ +1 ∥ ∥ ∥ . Therefore we know

This suggests that all gradients with respect to the weights become scaled by a factor of λ L -1 . The gradients with respect to the biases are proportional to λ L -l . When λ is large enough to render the gradient of the bias term negligible, it follows that ∥ ∥ ∇ θ f θ ′ ( x ) -λ L -1 ∇ θ f θ ( x ) ∥ ∥ = o (1) · ∥∇ θ f θ ′ ( x ) ∥ . 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 is the inner product between gradients, we know G θ ′ ( x , x ′ ) 's relative error to λ 2( L -1) G θ ( x , x ′ ) is o (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 1 Suppose that the network's parameter at iteration t is θ t . For each transition ( s i , a i , s i +1 , r i ) in dataset, denote r = [ r 1 , . . . , r M ] ⊤ ∈ R M , ˆ π θ t ( s ) = arg max a ˆ Q θ t ( s, a ) . Denote x ∗ i,t = ( s i +1 , ˆ π θ t ( s i +1 )) . Concatenate all x ∗ i,t to be X ∗ t . Denote u t = f θ t ( X ) -( r + γ · f θ t ( X ∗ t )) to be TD error vector at iteration t . The learning rate η is infinitesimal. When the maximal point of ˆ Q θ t is stable as t increases, we have the following evolving equation for u t +1

Proof. For the sake of simplicity, denote Z t = ∇ θ f θ ( X ) ∣ ∣ ∣ θ t , Z ∗ t = ∇ θ f θ ( X ∗ t ) ∣ ∣ ∣ θ t . The Q-value iteration minimizes loss function L defined by L ( θ ) = 1 2 ∥ f θ ( X ) -( r + γ · f θ t ( X ∗ t )) ∥ 2 2 . Therefore we have the gradient as

According to gradient descent, we know θ t +1 = θ t -η Z t u t . Since η is very small, we know θ t +1 stays within the neighborhood of θ t . We can Taylor-expand function f θ ( · ) near θ t as

When η is infinitesimally small, the equation holds. Plug in θ t +1 , we know

According to Assumption 2, we know ∥ X ∗ t -X ∗ t +1 ∥ = o ( η ) . So we can substitute X ∗ t +1 into X ∗ t since it is high order infinitesimal. Therefore, u t +1 boils down to

Theorem 3 Suppose we use SGD optimizer for Q-value iteration with learning rate η to be infinitesimal. Given iteration t > t 0 , and A = γ ¯ G ( ¯ X ∗ , X ) -¯ G ( X , X ) , where ¯ G is the Gram matrix under terminal kernel ¯ k . The divergence of u t is equivalent to whether there exists an eigenvalue λ of A such that Re( λ ) > 0 . If converge, we have u t = ( I + η A ) t -t 0 · u t 0 . Otherwise, u t becomes parallel to the eigenvector of the largest eigenvalue λ of A , and its norm diverges to infinity at following order

for some constant C ′ to be determined and L is the number of layers of MLP. Specially, when L = 2 , it reduces to O ( 1 1 -C ′ ληt ) .

Proof . According to Assumption 2, max action becomes stable after t 0 . It implies ∥ X ∗ t -¯ X ∗ ∥ = o (1) . The stability of the NTK direction implies that for some scalar k t and the specific input X ∗ , X , we have G θ t ( ¯ X ∗ , X ) = k t ¯ G ( X ∗ , X ) + ϵ t, 1 and G θ t ( X , X ) = k t ¯ G ( X , X ) + ϵ t, 2 where ∥ ϵ t,i ∥ = o ( k t ) , i = 1 , 2 . Therefore, we have

k t equals 1 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 δ t = γ ϵ t, 3 + ϵ t, 1 -ϵ t, 2 is of o (1) norm and thus negligible, we know all the eigenvalues of I + η A have form 1+ ηλ i . Considering η is small enough, we have | 1 + ηλ i | 2 ≈ 1 + 2 η Re( λ ) . Now suppose if there does not exists eigenvalue λ of A satisfies Re( λ ) > 0 , we have | 1 + ηλ i | ≤ 1 . Therefore, the NTK will become perfectly stable so k t = 1 for t > t 0 , and we have

Otherwise, there exists an eigenvalue for A satisfying Re( λ ) > 0 . Denote the one with the largest real part as λ , and v to be the corresponding eigenvector. We know matrix I + η A also has left eigenvector v , whose eigenvalue is 1 + ηλ . In this situation, we know after each iteration, ∥ u t +1 ∥ will become larger than ∥ u t ∥ . Moreover, to achieve larger and larger prediction values, the model's parameter's norm ∥ θ t ∥ also starts to explode. We know u t is homogeneous with respect θ t for ReLU networks. The output f θ t ( X ) enlarges p L times when θ t enlarges p times. When the reward values is small with respect to the divergent Q-value, TD error u t = O ( f θ t ( X )) = O ( θ L t ) . Besides, according to Lemma1, we know k t = O ( ∥ θ t ∥ 2( L -1) ) = O ( ∥ u t ∥ 2( L -1) /L ) = O ( ∥ u t ∥ 2 -2 /L ) .

Denote g ( ηt ) = v ⊤ u t , left multiply v to equation u t +1 = ( I + ηk t A ) u t . we have g ( ηt + η ) = (1 + ηλk t ) g ( ηt ) . Since we know such iteration will let u t to be dominated by v and align with v , we know g ( ηt ) = O ( ∥ u t ∥ ) for large t . Therefore k t = O ( ∥ u t ∥ 2( L -1) /L ) = C · g ( ηt ) 2 -2 /L . This boils down to g ( ηt + η ) = g ( ηt ) + Cηλg ( ηt ) 2 , which further becomes

Let η → 0 , we have an differential equation d g d t = Cλg ( t ) 3 -2 /L . When L = 1 , the MLP network degenerates to a linear function. The solution of ODE is

reflecting the exponential growth under linear function that has been studied in previous works [47]. When L > 2 , Solving this ODE gives

So at an infinite limit, we know ∥ u t ∥ = g ( ηt ) = O ( 1 (1 -C ′ ληt ) L/ (2 L -2) ) . Specially, for the experimental case we study in Figure 2 where L = 2 , it reduces to O ( 1 1 -C ′ ληt ) . We conduct more experiments with L = 3 in Appendix C.2 to verify our theoretical findings.

It is also worthwhile to point out another important mechanism that we found when the Q-value network is in a steady state of diverging. Although we assume that the learning rate η is infinitesimal, the arg max function is not L -Lipshitz. So why could both NTK function and optimal policy action X ∗ t converge in direction? We found that when the model enters such 'terminal direction', ( s i , ˆ π θ t ( s i )) not only becomes the max value state-action pair, but also satisfies

This means each x ∗ t,i = ( s i , ˆ π θ t ( s i )) 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 parameter to evolve in static dynamics , which in turn reassures it to be policy action in the next iteration. Therefore, after the critical point t 0 , 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 a particular direction, leading to an infinite growth of the model's parameters in that same direction. This phenomenon is 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 1 after reaching a critical point, even as the norm of the parameters continually increases.

C.2 Terminal Time

Theorem 3 claims ∥ u t ∥ = O ( f θ t ( X )) = O ( 1 (1 -C ′ ληt ) L/ (2 L -2) ) , implying the relation

Beisdes, it implies the existence of a 'terminal time' 1 C ′ ηλ that the model must crash at a singular point. When the training approaches this singular point, the estimation value and the model's norm explode rapidly in very few steps. We have run an experiment with L = 2 in Figure 2, from which we can see that Q-value's inverse proves to decay linearly and eventually becomes Nan at the designated time step. When L = 3 , from our theretical analysis, we have 1 /q 4 3 ∝ 1 -C ′ ληt . 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 1 -C ′ ληt after a critical point and it eventually reaches a NAN value at the terminal time step.

400

600

800

1000

Steps

0.0

0.5

1.0

1.5

2.0

2.5

3.0

1/

q

4

3

y = -0.00270 * x + 2.83

R-squared = 0.99855

1/

q

4/3

v.s. step

C.3 Adam Case

In this section, we will prove that if the algorithm employs Adam as the optimizer, the model still suffers divergence. Moreover, we demonstrate that the norm of the network increase linearly, of which the slope is η √ P , where P is the number of parameters and η is the learning rate. Also, the Q-value prediction will increase at L th -polynomial's rate, where L is the number of layers of model f θ . 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 in Figure 16, Figure 17, and Figure 18.

Theorem 4. Suppose we use Adam optimizer for Q-value iteration and all other settings are the same as Theorem 3. After t > t 0 , the model will diverge if and only if λ max ( A ) > 0 . If it diverges, we have ∥ θ t ∥ = η √ Pt + o ( t ) and ∥ u t ∥ = Θ( t L ) where P and L are the number of parameters and the number of layers for network f θ , 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

Recall that each iteration of the Adam algorithm has the following steps.

Instead of exactly solving this series, we can verify linear growth is indeed the terminal behavior of θ t since we only care about asymptotic order. Assume that θ t = kt for t > T , we can calculate m t by dividing both sides of the definition of m t by β t 1 , which gives

, where g t is given in Equation (24). Similarly, we have

Hence we verify that

therefore we know each iteration will increase each parameter by exactly constant η . This in turn verified our assumption that parameter θ t grows linearly. The slope for the overall parameter is thus η √ P . This can also be verified in Figure 4. When we have θ t = η √ P ¯ θ , 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 O ( t L ) .

D LayerNorm's Effect on NTK

In this section, we demonstrate the effect of LayerNorm on SEEM. Our demonstration is just an intuitive explanation rather than a rigorous proof. We show that adding a LayerNorm can effectively reduce the NTK between any x 0 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 ∗ t , X ) 's eigenvalue are greatly reduced when every individual NTK value between any x 0 and extreme input x is reduced.

We consider a two-layer MLP. The input is x ∈ R d in , and the hidden dimension is d . The parameters include W = [ w 1 , . . . , w d ] ⊤ ∈ R d × d in , b ∈ R d and a ∈ R d . Since for the NTK value, the last layer's bias term has a constant gradient, we do not need to consider it. The forward function of the network is

Proposition 1. For any input x and network parameter θ , if ∇ θ f θ ( x ) = 0 , then we have

̸

Similarly, we have

So we have

Denote N = { 1 , . . . , d } \ P . We know for every j ∈ N either ∂f ( x ) ∂a j = ∂f ( x ) ∂b j = ∂f ( x ) ∂ w j = 0 , or w ⊤ j x < 0 . For the latter case, we know lim λ →∞ ∂f ( λ x ) ∂a j = ∂f ( λ x ) ∂b j = ∂f ( λ x ) ∂ w j = 0 . In both cases, we have

Therefore, according to the definition of NTK, we have

For the model equipped with LayerNorm, the forward function becomes

where ψ ( · ) is the layer normalization function defined as

Denote P = I -11 ⊤ /d , note that the derivative of ψ ( · ) is

Now we state the second proposition.

Proposition 2. For any input x and network parameter θ and any direction v ∈ R d in , if the network has LayerNorm, then we know there exists a universal constant C , such that for any λ ≥ 0 , we have

̸

Proof. Denote z i = w ⊤ i x + b i , according to condition ∇ θ f θ ( x ) = 0 , we know there must exist at least one i such that z i > 0 , denote this set as P . Now consider all the i ∈ [ d ] that satisfy z i > 0 and w ⊤ i x > 0 (otherwise take opposite sign of λ ), we have

Specially, we have

Proof. Since for finite range, there always exists a constant upper bound, we just need to analyze the case for λ → + ∞ and shows that it is constant bounded. First compute ∇ θ f θ ( x ) and get

These quantities are all constant bounded. Next we compute lim λ →∞ ∇ θ f θ ( x + λ v )

According to the property of LayerNorm in Equation (43), we have

Therefore we know its limit is also constant bounded. So we know there exists a universal constant with respect to θ , x , v such that k NTK ( x , x + λ v ) = 〈 ∂f ( x ) ∂ θ i , ∂f ( x + λ v ) ∂ θ i 〉 ≤ C.

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, β 1 = 0 . 9 , and β 2 = 0 . 999 . The discount factor, γ , was set to 0.99. Our code builds upon the existing TD3+BC framework, which can be found at https://github.com/sfujim/TD3_BC .

SEEMReduction Experiments For the experiments discussed in Section Section 4, we maintained the same configuration as in 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.

Offline 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, DiffQL, and CQL as baselines. We implement our method on the top of official implementations of TD3+BC and IQL; for CQL, we use a reliable JAX implementation at https://github.com/ sail-sg/offbench [25]; for Diffusion Q-learning ( i . e ., combining TD3+BC and diffusion policy), we use an 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 0. Considering the extreme suboptimality in Antmaze medium and large datasets, we adopt Return-based Offline Prioritized Replay (OPER-R) with p base = 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 1 1 -C ′ ληt and is difficult to capture, we opted to use a simple toy task for these experiments. The task includes a continuous two-dimensional state space s = ( x 1 , x 2 ) ∈ S = R 2 , 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 a 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, we randomly sampled 100 state-action pairs as our offline dataset. The Q-network was implemented as a two-layer MLP with a hidden size of 200. We used SGD with a learning rate of 0.01, and the discount factor, γ was set to 0.99.

0

200

400

600

800

1000

K steps

0

1000

2000

3000

4000

5000

Average Return

walker2d

SAC w. EMA

SAC w.o. EMA

SAC w.o. EMA (LayerNorm)

0

200

400

600

800

1000

K steps

0

1000

2000

3000

Average Return

hopper

F More Experiments

F.1 Validate the homogeneity in Lemma 1.

We define a family of 3-layer ReLU-activated MLPs (with bias term) with different scaling λ of the same network parameter θ (from D4RL training checkpoints when value divergence happens), and

feed these networks with the same input. We show one example in the below table, which confirms that the L -degree homogeneity of output, gradient and NTK is valid at high precision, the NTK is almost parallel, too. The larger the scaling factor is, the more accurate k L increasing pattern is for output ( k L -1 for gradient and k 2( L -1) for NTK). Further, we empirically validated the L -degree homogeneity of NTK holds for all checkpoints in D4RL experiments where divergence happens.

F.2 Ablate the contributions of EMA, Double-Q, and LayerNorm.

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 in Figure 12. It 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.

50% Mujoco Dataset

10% Mujoco Dataset

0

200

400

600

Performance

w.o. EMA

+EMA

+DoubleQ

+LayerNorm

F.3 Benchmarking Normalizations.

Previously, we have demonstrated that LayerNorm, BatchNorm, and WeightNorm can effectively maintain a low SEEM and stabilize Q convergence in Section 4. Our next goal is to identify the most suitable regularization method for the value network in offline RL. Prior research has shown that divergence is correlated with poor control performance[48; 19]. In this context, we 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 or 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 is 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 a suitable normalization for the critic in offline RL.

TD3+BC

LN

BN

WN

300

350

400

450

500

550

600

Performance

Regularizations

F.4 How LayerNorm should be added.

The inclusion of LayerNorm is situated between the linear and activation layers. However, the ideal configuration for adding LayerNorm can vary and may depend on factors such as 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, we conducted an assessment of their impacts on performance in the two most challenging Antmaze environments. Our experimental setup mirrored that of the Antmaze experiments mentioned above, utilizing a three-layer MLP critic with a hidden size configuration of (256,256,256). We evaluated variants where LayerNorm was only applied to a portion of hidden layers and where learnable affine parameters were disabled. As seen in Table 5, the performances with LayerNorm applied solely to the initial layers LN (0), LN (0,1) are considerably lower compared to the other setups in the 'antmaze-large-play-v0' task, while applying LayerNorm to all layers LN(0,1,2) seems to yield the best performance. For the 'antmaze-large-diverse-v0' task, performances seem to 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.

F.5 Baird's Counterexample.

Here, we will present our solution allow linear approximation to 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 of linear approximation could effectively lead to convergence (see Figure 14). 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.

G Discussion

SEEM and Deadly Triad. Deadly Triad is a term that refers to a problematic interaction observed in reinforcement learning algorithms, where off-policy learning, function approximation, and bootstrapping 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 in 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 ∗ t is merely a perturbation of X . Consequently, A t = γ G θ t ( X ∗ t , X ) -G θ t ( X , X ) ≈ ( γ -1) G θ t ( X , X ) , with G tending to be negative-definite. Without function approximation, the update of Q ( X ) will not influence Q ( X ∗ t ) , and the first term in A t becomes zero. A t = -G θ t ( X , X ) ensures that SEEM is non-positive 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 are present, the NTK analysis gives rise to the form A t = γ G θ t ( X ∗ t , X ) -G θ t ( 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), ReLUactivated 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 λx , the activation state for every neuron in the network becomes deterministic when λ →∞ . Therefore, the ReLU activation layer becomes a multiplication of a constant 0-1 mask, and the whole network degenerates to a linear function. If the network becomes a linear function for inputs λx 0 , λ ≥ C , this will cause an arbitrarily large NTK value between extreme points λx 0 and dataset point x 0 . Because now f θ ( λx ) can be written in equivalent form W T ( λx ) , so ϕ ( λx ) = λW becomes linear proportional to λ . This further induces linear NTK value between x and λx that is unbounded. This phenomenon can be intuitively interpreted as below, since the value prediction of f θ ( x ) on line λx 0 is almost a linear function W T x for large λ , any subtle change of effective parameter W will induce a large change on λx 0 for large λ . This demonstrates a strong correlation between the value prediction between point x 0 and far-away extreme point λx 0 . 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 a better balance.

0

250

500

750

1000

1250

1500

1750

2000

Steps

0

1000

2000

3000

4000

5000

6000

Weight Norm

Q-learning on Baird's Counterexample

No LayerNorm

LayerNorm

0

250

500

750

1000

K Steps

0.0

0.2

0.4

0.6

0.8

1.0

antmaze-large-diverse-v0

Policy Action

Data Action

0

250

500

750

1000

K Steps

0.0

0.2

0.4

0.6

0.8

1.0

antmaze-umaze-v0

0

250

500

750

1000

K Steps

0.0

0.2

0.4

0.6

0.8

1.0

antmaze-large-play-v0

0

250

500

750

1000

K Steps

0.0

0.2

0.4

0.6

0.8

1.0

antmaze-umaze-diverse-v0

0

250

500

750

1000

K Steps

0.0

0.2

0.4

0.6

0.8

1.0

hopper-medium-v2

0

250

500

750

1000

K Steps

0.0

0.2

0.4

0.6

0.8

1.0

walker2d-medium-v2

0

250

500

750

1000

K Steps

0.0

0.2

0.4

0.6

0.8

1.0

0

250

500

750

1000

K Steps

0.0

0.2

0.4

0.6

0.8

1.0

walker2d-medium-expert-v2

Action Norm / Extreme Norm

H More Visualization Results

In Assumption 2, we posit that the direction of NTK and the policy remains stable following a 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 1 as shown in 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 in the estimated Q-value.

0

1000

2000

3000

K steps

0.0

0.2

0.4

0.6

0.8

1.0

NTK similarity

0

1000

2000

3000

K steps

0.0

0.2

0.4

0.6

0.8

1.0

Action similarity

0

1000

2000

3000

K steps

0.0

0.2

0.4

0.6

0.8

1.0

Model similarity

0

1000

2000

3000

K steps

0

50000

100000

150000

200000

walker2d-medium-v2 (3 seeds)

0

1000

2000

3000

K steps

0.0

0.2

0.4

0.6

0.8

1.0

NTK similarity

0

1000

2000

3000

K steps

0.0

0.2

0.4

0.6

0.8

1.0

Action similarity

0

1000

2000

3000

K steps

0.0

0.2

0.4

0.6

0.8

1.0

Model similarity

0

1000

2000

3000

K steps

0

50000

100000

150000

200000

walker2d-medium-replay-v2 (3 seeds)

0

1000

2000

3000

K steps

0.0

0.2

0.4

0.6

0.8

1.0

NTK similarity

0

1000

2000

3000

K steps

0.0

0.2

0.4

0.6

0.8

1.0

Action similarity

0

1000

2000

3000

K steps

0.0

0.2

0.4

0.6

0.8

1.0

Model similarity

0

1000

2000

3000

K steps

0

50000

100000

150000

200000

walker2d-medium-expert-v2 (3 seeds)

0

1000

2000

3000

K steps

0.0

0.2

0.4

0.6

0.8

1.0

NTK similarity

0

1000

2000

3000

K steps

0.0

0.2

0.4

0.6

0.8

1.0

Action similarity

0

1000

2000

3000

K steps

0.0

0.2

0.4

0.6

0.8

1.0

Model similarity

0

1000

2000

3000

K steps

0

50000

100000

150000

200000

0

1000

2000

3000

K steps

0.0

0.2

0.4

0.6

0.8

1.0

NTK similarity

0

1000

2000

3000

K steps

0.0

0.2

0.4

0.6

0.8

1.0

Action similarity

0

1000

2000

3000

K steps

0.0

0.2

0.4

0.6

0.8

1.0

Model similarity

0

1000

2000

3000

K steps

0

50000

100000

150000

0

1000

2000

3000

K steps

0.0

0.2

0.4

0.6

0.8

1.0

NTK similarity

0

1000

2000

3000

K steps

0.0

0.2

0.4

0.6

0.8

1.0

Action similarity

0

1000

2000

3000

K steps

0.0

0.2

0.4

0.6

0.8

1.0

Model similarity

0

1000

2000

3000

K steps

0

50000

100000

150000

0

1000

2000

3000

K steps

0.0

0.2

0.4

0.6

0.8

1.0

NTK similarity

0

1000

2000

3000

K steps

0.0

0.2

0.4

0.6

0.8

1.0

Action similarity

0

1000

2000

3000

K steps

0.0

0.2

0.4

0.6

0.8

1.0

Model similarity

0

1000

2000

3000

K steps

0

50000

100000

150000

200000

hopper-medium-v2 (3 seeds)

0

1000

2000

3000

K steps

0.0

0.2

0.4

0.6

0.8

1.0

NTK similarity

0

1000

2000

3000

K steps

0.0

0.2

0.4

0.6

0.8

1.0

Action similarity

0

1000

2000

3000

K steps

0.0

0.2

0.4

0.6

0.8

1.0

Model similarity

0

1000

2000

3000

K steps

0

50000

100000

150000

200000

hopper-medium-replay-v2 (3 seeds)

0

1000

2000

3000

K steps

0.0

0.2

0.4

0.6

0.8

1.0

NTK similarity

0

1000

2000

3000

K steps

0.0

0.2

0.4

0.6

0.8

1.0

Action similarity

0

1000

2000

3000

K steps

0.0

0.2

0.4

0.6

0.8

1.0

Model similarity

0

1000

2000

3000

K steps

0

50000

100000

150000

hopper-medium-expert-v2 (3 seeds)

0

1000

2000

3000

K steps

0.0

0.2

0.4

0.6

0.8

1.0

NTK similarity

0

1000

2000

3000

K steps

0.0

0.2

0.4

0.6

0.8

1.0

Action similarity

0

1000

2000

3000

K steps

0.0

0.2

0.4

0.6

0.8

1.0

Model similarity

0

1000

2000

3000

K steps

0

50000

100000

150000

200000

antmaze-umaze-v0 (3 seeds)

0

1000

2000

3000

K steps

0.0

0.2

0.4

0.6

0.8

1.0

NTK similarity

0

1000

2000

3000

K steps

0.0

0.2

0.4

0.6

0.8

1.0

Action similarity

0

1000

2000

3000

K steps

0.0

0.2

0.4

0.6

0.8

1.0

Model similarity

0

1000

2000

3000

K steps

0

50000

100000

150000

200000

antmaze-umaze-diverse-v0 (3 seeds)

0

1000

2000

3000

K steps

0.0

0.2

0.4

0.6

0.8

1.0

NTK similarity

0

1000

2000

3000

K steps

0.0

0.2

0.4

0.6

0.8

1.0

Action similarity

0

1000

2000

3000

K steps

0.0

0.2

0.4

0.6

0.8

1.0

Model similarity

0

1000

2000

3000

K steps

0

50000

100000

150000

200000

0

1000

2000

3000

K steps

0.0

0.2

0.4

0.6

0.8

1.0

NTK similarity

0

1000

2000

3000

K steps

0.0

0.2

0.4

0.6

0.8

1.0

Action similarity

0

1000

2000

3000

K steps

0.0

0.2

0.4

0.6

0.8

1.0

Model similarity

0

1000

2000

3000

K steps

0

10000

20000

30000

0

1000

2000

3000

K steps

0.0

0.2

0.4

0.6

0.8

1.0

NTK similarity

0

1000

2000

3000

K steps

0.0

0.2

0.4

0.6

0.8

1.0

Action similarity

0

1000

2000

3000

K steps

0.0

0.2

0.4

0.6

0.8

1.0

Model similarity

0

1000

2000

3000

K steps

0

25000

50000

75000

100000

125000

antmaze-large-play-v0 (3 seeds)

0

1000

2000

3000

K steps

0.0

0.2

0.4

0.6

0.8

1.0

NTK similarity

0

1000

2000

3000

K steps

0.0

0.2

0.4

0.6

0.8

1.0

Action similarity

0

1000

2000

3000

K steps

0.0

0.2

0.4

0.6

0.8

1.0

Model similarity

0

1000

2000

3000

K steps

0

25000

50000

75000

100000

125000

antmaze-large-diverse-v0 (3 seeds)

0

2000

K steps

0.0

0.5

1.0

SEEM/Q-value

0

2000

K steps

0.0

0.5

1.0

SEEM/Q-value

0

2000

K steps

0.0

0.5

1.0

SEEM/Q-value

10

10

10

13

10

10

10

13

10

10

10

13

walker2d-medium-v2 (3 seeds)

0

2000

K steps

0.0

0.5

1.0

SEEM/Q-value

0

2000

K steps

0.0

0.5

1.0

SEEM/Q-value

0

2000

K steps

0.0

0.5

1.0

SEEM/Q-value

10

10

10

13

10

10

10

13

10

10

10

13

walker2d-medium-replay-v2 (3 seeds)

0

2000

K steps

0.0

0.5

1.0

SEEM/Q-value

0

2000

K steps

0.0

0.5

1.0

SEEM/Q-value

0

2000

K steps

0.0

0.5

1.0

SEEM/Q-value

10

10

10

13

10

10

10

13

10

10

10

13

walker2d-medium-expert-v2 (3 seeds)

0

2000

K steps

0.0

0.5

1.0

SEEM/Q-value

0

2000

K steps

0.0

0.5

1.0

SEEM/Q-value

0

2000

K steps

0.0

0.5

1.0

SEEM/Q-value

10

10

10

13

10

10

10

13

10

10

10

13

0

2000

K steps

0.0

0.5

1.0

SEEM/Q-value

0

2000

K steps

0.0

0.5

1.0

SEEM/Q-value

0

2000

K steps

0.0

0.5

1.0

SEEM/Q-value

10

9

10

12

10

8

10

12

10

10

10

13

0

2000

K steps

0.0

0.5

1.0

SEEM/Q-value

0

2000

K steps

0.0

0.5

1.0

SEEM/Q-value

0

2000

K steps

0.0

0.5

1.0

SEEM/Q-value

10

10

10

13

10

10

10

13

10

10

10

13

0

2000

K steps

0.0

0.5

1.0

SEEM/Q-value

0

2000

K steps

0.0

0.5

1.0

SEEM/Q-value

0

2000

K steps

0.0

0.5

1.0

SEEM/Q-value

10

9

10

12

10

9

10

12

10

9

10

12

hopper-medium-v2 (3 seeds)

0

2000

K steps

0.0

0.5

1.0

SEEM/Q-value

0

2000

K steps

0.0

0.5

1.0

SEEM/Q-value

0

2000

K steps

0.0

0.5

1.0

SEEM/Q-value

10

9

10

12

10

9

10

12

10

9

10

12

hopper-medium-replay-v2 (3 seeds)

0

2000

K steps

0.0

0.5

1.0

SEEM/Q-value

0

2000

K steps

0.0

0.5

1.0

SEEM/Q-value

0

2000

K steps

0.0

0.5

1.0

SEEM/Q-value

10

9

10

12

10

9

10

12

10

9

10

12

hopper-medium-expert-v2 (3 seeds)

0

1000

2000

3000

K steps

0.0

0.5

1.0

SEEM/Q-value

0

1000

2000

3000

K steps

0.0

0.5

1.0

SEEM/Q-value

0

1000

2000

3000

K steps

0.0

0.5

1.0

SEEM/Q-value

10

4

10

10

10

9

10

12

10

4

10

10

antmaze-umaze-v0 (3 seeds)

0

1000

2000

3000

K steps

0.0

0.5

1.0

SEEM/Q-value

0

1000

2000

3000

K steps

0.0

0.5

1.0

SEEM/Q-value

0

1000

2000

3000

K steps

0.0

0.5

1.0

SEEM/Q-value

10

4

10

10

10

2

10

4

10

8

10

11

antmaze-umaze-diverse-v0 (3 seeds)

0

1000

2000

3000

K steps

0.0

0.5

1.0

SEEM/Q-value

0

1000

2000

3000

K steps

0.0

0.5

1.0

SEEM/Q-value

0

1000

2000

3000

K steps

0.0

0.5

1.0

SEEM/Q-value

10

4

10

10

10

7

10

10

10

4

10

10

0

1000

2000

3000

K steps

0.0

0.5

1.0

SEEM/Q-value

0

1000

2000

3000

K steps

0.0

0.5

1.0

SEEM/Q-value

0

1000

2000

3000

K steps

0.0

0.5

1.0

SEEM/Q-value

10

6

10

9

10

4

10

8

10

2

10

6

0

1000

2000

3000

K steps

0.0

0.5

1.0

SEEM/Q-value

0

1000

2000

3000

K steps

0.0

0.5

1.0

SEEM/Q-value

0

1000

2000

3000

K steps

0.0

0.5

1.0

SEEM/Q-value

10

3

10

9

10

3

10

1

10

3

10

9

antmaze-large-play-v0 (3 seeds)

0

1000

2000

3000

K steps

0.0

0.5

1.0

SEEM/Q-value

0

1000

2000

3000

K steps

0.0

0.5

1.0

SEEM/Q-value

0

1000

2000

3000

K steps

0.0

0.5

1.0

SEEM/Q-value

10

3

10

9

10

2

10

1

10

0

10

4

10

10

antmaze-large-diverse-v0 (3 seeds)

0

2000

K steps

10

3

10

12

Unnormalized SEEM

0

2000

K steps

10

11

10

15

Unnormalized SEEM

0

2000

K steps

10

3

10

11

Unnormalized SEEM

antmaze-umaze-v0 (3 seeds)

0

2000

K steps

10

3

10

11

Unnormalized SEEM

0

2000

K steps

10

1

10

2

Unnormalized SEEM

0

2000

K steps

10

4

10

11

Unnormalized SEEM

antmaze-umaze-diverse-v0 (3 seeds)

0

2000

K steps

10

3

10

11

Unnormalized SEEM

0

2000

K steps

10

6

10

10

Unnormalized SEEM

0

2000

K steps

10

3

10

12

Unnormalized SEEM

0

2000

K steps

10

2

10

7

Unnormalized SEEM

0

2000

K steps

10

2

10

8

Unnormalized SEEM

0

2000

K steps

10

1

10

4

Unnormalized SEEM

0

2000

K steps

10

3

10

12

Unnormalized SEEM

0

2000

K steps

10

4

10

2

Unnormalized SEEM

0

2000

K steps

10

2

10

10

Unnormalized SEEM

antmaze-large-play-v0 (3 seeds)

0

2000

K steps

10

3

10

12

Unnormalized SEEM

0

2000

K steps

10

4

10

2

Unnormalized SEEM

0

2000

K steps

10

3

10

11

Unnormalized SEEM

antmaze-large-diverse-v0 (3 seeds)

Figure 20: In Figure (a), an inflation in the estimated Q-value coincides with a surge in the normalized SEEM. However, there are some anomalies, such as 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 an 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.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n- 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.\n\nWEAKNESSES:\n- 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.\n\nQUESTIONS:\n1. 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\n\nLIMITATIONS:\n1. 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.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-06-07"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\n- 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.\n\nWEAKNESSES:\n- 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.\n\nQUESTIONS:\nPlease see the weaknesses above.\n\nLIMITATIONS:\nN/A", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-06-14"}

comments: SUMMARY:
This paper analyzes Q-value divergence in Offline-RL by considering a
neural tangent kernel for the value function.  They show that
consideration of this kernel is predictive of Q-value divergence.
This analysis further leads to the observation that using a LayerNorm
yields a kernel that behaves more like one would hope -- with nearer
values being more impacted than farther ones.  Their empirical tests
on D4RL benchmarks similarly show the benefits of LayerNorm.

STRENGTHS:
I like the insights in this paper and suspect that its publication
will lead others to further engineering that will yield even better
performing adjustments than the simple addition of the LayerNorm.
Figure 5 suggests there is more to be done.  Why does the figure on
the left indicate almost the opposite of what one would hope for.  At
a minimum, should the value right at x_0 be red?

WEAKNESSES:
The paper needs a number of editing improvements.  First, it needs an
English grammar checker.  The errors mostly don't lead to difficulties
in understanding, but, the paper should not be published with the
current level of English quality.

In section 5.1, the BC term should be explained even though the
reference is given for it.

The reference to a score of 81.4 in that section does not appear to match
the numbers in Table 1.

The legends in the figures need to use a larger font size.

QUESTIONS:
None

LIMITATIONS:
N/A
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 4
DATE: 2023-07-05