SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Yang Yue\*1 Rui Lu\*1 Bingyi Kang\*2 Shiji Song1 Gao Huang\*1
$^{1}$  Department of Automation, BNRist, Tsinghua University  $^{2}$  ByteDance Inc.
{le-y22, r-lu21}@ mails.tsinghua.edu.cn
bingykang@gmail.com{shijis,gaohuang}@tsinghua.edu.cn
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 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 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.
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.
How does divergence actually occur? Most existing studies only provide analyses of linear functions as Q-values, which is usually confined to contrived toy examples [47; 4]. The understanding and explanation for the divergence of non-linear neural networks in practice are still lacking. Little is known about what transpires in the model when its estimation inflates to infinity and what essential mechanisms behind the deadly triad contribute to this phenomenon. We are going to answer this question from the perspective of function approximation in the context of offline RL.
How to avoid detrimental bias? Despite the effectiveness of policy constraint methods in offline RL, they pose potential problems by introducing detrimental bias. Firstly, they explicitly force the policy to be near the behavior policy, which can potentially hurt performance. Also, the trade-off between performance and constraint needs to be balanced manually for each task. Secondly, these methods become incapable of dealing with divergence when 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 Self-Excite Eigenvalue Measure (SEEM), which is defined as the largest eigenvalue in the linear iteration during training, serving as an indicator of divergence. It also elucidates the fundamental mechanisms driving this divergence.
- 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.
Notation. A Markov Decision Process (MDP) is defined by tuple  $\mathcal{M} = (S, A, P, r, \nu)$ , where  $S \subseteq \mathbb{R}^{d_S}$  is the state space and  $A \subseteq \mathbb{R}^{d_A}$  is the action space.  $P: S \times A \mapsto \Delta(S)$  is the transition dynamics mapping from state-action pair to distribution in state space.  $r: S \times A \mapsto \mathbb{R}$  is the reward function and  $\nu \in \Delta(S)$  is the initial state distribution. An offline RL algorithm uses an offline dataset
Figure 1: NTK similarity and action similarity along the training in two offline D4RL tasks. Different colors represent different seeds. The cosine similarity between the current step and the last step is computed. We can see all runnings reach a steady NTK direction and policy actions after a time  $t_0$ .
$D = \{(s_i, a_i, s_i', r_i)\}_{i=1}^M$ , which consists of  $M$  finite samples from interactions with MDP  $\mathcal{M}$ . We write  $s \in D$  if there exists some  $k \in [M]$  such that  $s_k = s$ , similar for notation  $(s, a) \in D$ .
In this paper, we consider using a  $L$ -layer ReLU-activated MLP as the Q-value approximator. Each layer has parameter weight  $\pmb{W}_{\ell} \in \mathbb{R}^{d_{\ell + 1} \times d_{\ell}}$  and bias  $\pmb{b}_{\ell} \in \mathbb{R}^{d_{\ell + 1}}$ . The activation function is ReLU:  $\sigma(x) = \max(x, 0)$ . Integer  $d_{\ell}$  represents the dimensionality of the  $\ell_{th}$  hidden layer, where  $d_0 = d_S + d_A$  and  $d_L = 1$  since a scalar estimation is required for Q-value. Denote  $\pmb{\theta} = \mathrm{Vec}[W_0, \pmb{b}_0, \dots, \pmb{w}_L, \pmb{b}_L]$  to be the vector of all learnable parameters. We use  $\hat{Q}_{\pmb{\theta}}(s, a) = f_{\pmb{\theta}}(s, a)$  to denote the estimated Q-value using neural network  $f_{\pmb{\theta}}$  with parameter  $\pmb{\theta}$ . Also,  $\hat{\pi}_{\pmb{\theta}}(s) = \arg \max_a \hat{Q}_{\pmb{\theta}}(s, a)$  is its induced policy. Sometimes we will write  $f_{\pmb{\theta}}(\pmb{X}) \in \mathbb{R}^M$  where  $\pmb{X} = [(s_1, a_1), \dots, (s_M, a_M)]$  stands for the concatenation of all inputs  $(s_i, a_i) \in D$  in dataset.
Considering a Q-learning-based offline RL algorithm that alternates between the following two steps:
$$
\bar {Q} \leftarrow r (s, a) + \gamma \max  _ {a ^ {\prime}} \hat {Q} _ {\boldsymbol {\theta}} \left(s ^ {\prime}, a ^ {\prime}\right), \quad \boldsymbol {\theta} \leftarrow \boldsymbol {\theta} - \eta \nabla_ {\boldsymbol {\theta}} \left(\bar {Q} - \hat {Q} _ {\boldsymbol {\theta}} (s, a)\right) ^ {2}, \tag {1}
$$
where  $\eta$  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  $\pmb{x}$  and  $\pmb{x}'$ , with dependency on parameter  $\pmb{\theta}$ . Given a neural network  $f_{\theta}$  and two points  $\pmb{x}, \pmb{x}'$  in the input space, the NTK is defined as  $k_{\mathrm{NTK}}(\pmb{x}, \pmb{x}') \coloneqq \langle \nabla_{\theta}f_{\theta}(\pmb{x}), \nabla_{\theta}f_{\theta}(\pmb{x}') \rangle$ , where the feature extraction is  $\phi_{\theta}(\pmb{x}) \coloneqq \nabla_{\theta}f_{\theta}(\pmb{x})$ . If we have two batches of input  $\pmb{X}, \pmb{X}'$ , we can compute the Gram matrix as  $G_{\theta}(\pmb{X}, \pmb{X}') = \phi_{\theta}(\pmb{X})^\top \phi_{\theta}(\pmb{X}')$ .
Intuitively speaking, NTK measures the similarity between two inputs through the lens of a non-linear network  $f_{\theta}$ . Such similarity can be simply understood as the correlation between the network's value predictions of these two inputs. If  $k_{\mathrm{NTK}}(\pmb{x},\pmb{x}^{\prime})$  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  $\pmb{\theta}$ , say, from  $\pmb{\theta}$  to  $\pmb{\theta}' = \pmb{\theta} + \Delta \pmb{\theta}$ , denote the prediction change for  $\pmb{x}$  as  $\Delta f(\pmb{x}) = f_{\pmb{\theta}'}(\pmb{x}) - f_{\pmb{\theta}}(\pmb{x})$  and for  $x'$  as  $\Delta f(x') = f_{\theta'}(x') - f_{\theta}(x')$ .  $\Delta f(\pmb{x})$  and  $\Delta f(x')$  will be highly correlated, and the strength of the correlation is proportional to  $k_{\mathrm{NTK}}(\pmb{x},\pmb{x}')$ . This intuition can be explained from simple Taylor expansion, which gives  $\Delta f(\pmb{x}) = \Delta \pmb{\theta}^\top \phi_{\pmb{\theta}}(\pmb{x})$  and  $\Delta f(x') = \Delta \pmb{\theta}^\top \phi_{\pmb{\theta}}(\pmb{x}')$ . These two quantities will be highly correlated if  $\phi_{\pmb{\theta}}(\pmb{x})$  and  $\phi_{\pmb{\theta}}(\pmb{x}')$  are well-aligned, which corresponds to large NTK value between  $\pmb{x}$  and  $x'$ . To summary,  $k_{\mathrm{NTK}}(\pmb{x},\pmb{x}')$  characterizes the subtle connection between  $f_{\theta}(\pmb{x})$  and  $f_{\theta}(x')$  caused by neural network's generalization.
In this section, we will investigate the divergence phenomenon of Q-value estimation and theoretically explain 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 valuable 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  $\pmb{\theta}_t$ . For each transition  $(s_i, a_i, s_{i+1}, r_i)$  in dataset, denote  $\pmb{r} = [r_1, \dots, r_M]^\top \in \mathbb{R}^M$ ,  $\hat{\pi}_{\pmb{\theta}_t}(s) = \arg \max_a \hat{Q}_{\pmb{\theta}_t}(s, a)$ . Denote  $\pmb{x}_{i,t}^* = (s_{i+1}, \hat{\pi}_{\pmb{\theta}_t}(s_{i+1}))$ . Concatenate all  $\pmb{x}_{i,t}^*$  to be  $\pmb{X}_t^*$ . Denote  $\pmb{u}_t = f_{\pmb{\theta}_t}(\pmb{X}) - (\pmb{r} + \gamma \cdot f_{\pmb{\theta}_t}(\pmb{X}_t^*))$  to be TD error vector at iteration  $t$ . The learning rate  $\eta$  is infinitesimal. When the maximal point of  $\hat{Q}_{\pmb{\theta}_t}$  is stable as  $t$  increases, we have the evolving equation for  $\pmb{u}_{t+1}$  as  $\pmb{u}_{t+1} = (\pmb{I} + \eta \pmb{A}_t)\pmb{u}_t$ , where  $\pmb{A}_t = (\gamma \phi_{\pmb{\theta}_t}(\pmb{X}_t^*) - \phi_{\pmb{\theta}_t}(\pmb{X}))^\top \phi_{\pmb{\theta}_t}(\pmb{X}) = \gamma G_{\pmb{\theta}_t}(\pmb{X}_t^*, \pmb{X}) - G_{\pmb{\theta}_t}(\pmb{X}, \pmb{X})$ .
$\pmb{G}$  is defined in Section 2. Detailed proof is left in Appendix B. Theorem 1 states that when the policy action  $\hat{\pi}_{\pmb{\theta}_t}(s)$  that maximizes Q-values keep stable, leading to invariance of  $\pmb{X}_t^*$ , each Q-value iteration essentially updates the TD error vector by a matrix  $\pmb{A}_t$ , which is determined by a discount factor  $\gamma$  and the NTK between  $\pmb{X}$  and  $\pmb{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  $\hat{\pi}_{\theta_t}(s)$  that maximizes Q-values may constantly fluctuate with  $\pmb{\theta}_t$  over the course of training, leading to variations of  $X_{t}^{*}$ . Also, the kernel  $G_{\theta}(\cdot ,\cdot)$  has dependency on parameter  $\pmb{\theta}$ . This causes non-linear dynamics of  $\pmb{u}_t$  and the model parameter vector  $\pmb{\theta}$ . 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  $\bar{k}(\cdot, \cdot)$  and stabilized policy state-action  $\bar{\mathbf{X}}^*$  such that  $\left\| \frac{k_t(\boldsymbol{x}, \boldsymbol{x}')}{\|\boldsymbol{x}\| \cdot \|\boldsymbol{x}'\|} - \frac{\bar{k}(\boldsymbol{x}, \boldsymbol{x}')}{\|\boldsymbol{x}\| \cdot \|\boldsymbol{x}'\|} \right\| = o(1)$  and  $\left\| \mathbf{X}_{t+1}^* - \mathbf{X}_t^* \right\| = o(\eta)$  for any  $\boldsymbol{x}, \boldsymbol{x}' \in S \times \mathcal{A}$  when  $t > t_0$ . We also assume that  $\left\| \mathbf{X}_t^* - \bar{\mathbf{X}}^* \right\| = o(1)$ .
This assumption states that the training dynamics will reach a steady state, where the policy  $\hat{\pi}_{\theta_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<sup>1</sup>, 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  $\hat{\pi}_{\theta_t}$  converges after a certain timestep  $t_0$ . Moreover, we find that  $\hat{\pi}_{\theta_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.
Figure 2: Prediction Ability - Linear decay of inverse Q-value with SGD and  $L = 2$
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  $\eta$  to be infinitesimal. Given iteration  $t > t_0$ , and  $\mathbf{A} = \gamma \bar{\mathbf{G}} (\bar{\mathbf{X}}^*,\mathbf{X}) - \bar{\mathbf{G}} (\mathbf{X},\mathbf{X})$ , where  $\bar{\mathbf{G}}$  is the Gram matrix under terminal kernel  $k$ . The divergence of  $\mathbf{u}_t$  is equivalent to whether there exists an eigenvalue  $\lambda$  of  $\mathbf{A}$  such that  $\mathrm{Re}(\lambda) > 0$ . If converge, we have  $\mathbf{u}_t = (\mathbf{I} + \eta \mathbf{A})^{t - t_0}\cdot \mathbf{u}_{t_0}$ . Otherwise,  $\mathbf{u}_t$  becomes parallel to the eigenvector of the largest eigenvalue  $\lambda$  of  $\mathbf{A}$ , and its norm diverges to infinity at following order
$$
\left\| \boldsymbol {u} _ {t} \right\| _ {2} = O \left(\frac {1}{\left(1 - C ^ {\prime} \lambda \eta t\right) ^ {L / \left(2 L - 2\right)}}\right) \tag {2}
$$
for some constant  $C'$  to be determined and  $L$  is the number of layers of MLP. Specially, when  $L = 2$ , it reduces to  $O\left(\frac{1}{1 - C' \lambda \eta t}\right)$ .
Theorem 3 predicts the divergence by whether  $\lambda_{\max}(A)$  is greater to 0. We term this divergence detector  $\lambda_{\max}(A)$  as Self-Excite Eigenvalue Measure, or SEEM. Essentially, Theorem 3 states that
To affirm the practicability of our analysis, we conduct experiments on a widely used offline RL benchmark D4RL [10], and utilize a 3-Layer MLP instead of a toy example. More details can be found in the appendix.
Figure 3: The divergence indication property of SEEM. The left figure shows the SEEM value with respect to different discount factors  $\gamma$ . Theorem 3 states that larger  $\gamma$  has a larger SEEM value. The star point means the model's prediction eventually exceeds a large threshold  $(10^{6})$  in training, which means diverging. We can see that positive SEEM is perfectly indicative of the divergence. This is also true within the training process. From the middle and right figures, we can see that the prediction Q-value (in blue) is stable until the normalized kernel matrix's SEEM (in red) rises up to a large positive value, then we can observe the divergence of the model.
we can monitor SEEM value to know whether the training will diverge. To validate this, we conduct experiments with different discount factor  $\gamma$ . The first term in  $A$  is linear to  $\gamma$ , thus larger  $\gamma$  is more susceptible to diverging. This is exactly what we find in reality. In Figure 3, we can see that experiments with large  $\gamma$  have positive SEEM, and the training eventually diverges ( $\star$ ). 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.
$\succ$  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 = \frac{1}{C' \lambda \eta}$ . 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 / \| \boldsymbol{u}_t \| \approx -1.44 \times 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.
$\succ$  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  $\lambda_{\max}(\mathbf{A}) > 0$ . If it diverges, we have  $\| \pmb{\theta}_t \| = \eta \sqrt{P} t + o(t)$  and  $\| \pmb{u}_t \| = \Theta (t^L)$  where  $P$  and  $L$  are the number of parameters and the number of layers for network  $f_{\pmb{\theta}}$ , 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  $\| \pmb{\theta}_t\|$  exhibits a straight line after a critic point  $t_0$ . Moreover, we can see  $\log Q \approx 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  $\lambda_{\max}(A)$  is as below: Think about only one sample  $x = (s_0, a_0)$  and its stabilized next step state-action
Figure 4: Prediction ability - Linear growth of network parameter's norm and polynomial growth of predicted Q-value with an Adam optimizer. Please note that in the 2th and 4th figures, both Q-value and steps have been taken the logarithm base 10. Different color represents different seeds.