SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$\pmb{x}^{*} = (s_{1},a_{1})$ . Without loss of generality, we assume  $f_{\theta_t}(\pmb {x}) < r + \gamma f_{\theta_{t + 1}}(\pmb{x}^*)$ . If SEEM is positive, we have  $\gamma G_{\theta_t}(\pmb{x}^*,\pmb {x})$  is larger than  $G_{\theta_t}(\pmb {x},\pmb {x})$ . Recall that  $G_{\theta_t}(\pmb{x}^*,\pmb {x})$  depicts the strength of the bond between  $\pmb{x}$  and  $\pmb{x}^*$  because of generalization. So we know that, when updating the value of  $f_{\theta}(\pmb {x})$  towards  $r + \gamma f_{\theta_{t + 1}}(\pmb{x}^*)$ , the Q-value iteration inadvertently makes  $f_{\theta_{t + 1}}(\pmb{x}^*)$  increase even more than the increment of  $f_{\theta_{t + 1}}(\pmb {x})$ . Consequently, the TD error  $r + \gamma f_{\theta}(\pmb{x}^{*}) - f_{\theta}(\pmb {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  $\hat{\pi}_{\theta_t}(s)$ . Previous Theorem 1 and Theorem 3 rely on the stability of the policy action  $\hat{\pi}_{\theta_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  $\hat{\pi}_{\theta_t}(s)$ . Consequently, as aforementioned, a loop is formed:  $\hat{Q}(s,a)$  keeps chasing  $\hat{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  $\mathbf{G}(\mathbf{X},\mathbf{X}) = \mathbf{X}^{\top}\mathbf{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(\theta) = 0$ , and it requires the label noise  $\varepsilon$  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  $\theta$  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  $\phi(s', a')$  to get  $\overline{\mathcal{R}}_{exp}(\theta)$ . 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  $\boldsymbol{x}_0 = (0.1, 0.2)$  and calculate the NTK value  $\boldsymbol{G}_{\theta}(\boldsymbol{x}_0, \boldsymbol{x})$  for a range of  $\boldsymbol{x}$  in the input domain. The heatmap displays high values at the boundaries of the input range,
Figure 5: The normalized NTK map for 2-layer ReLU MLP with and without layernorm. The input and hidden dimensions are 2 and 10,000. We can see that for the MLP without LayerNorm, prediction change at  $x_0$  has a dramatic influence on points far away like  $x = (4,4)$ , meaning a slight change of  $f(x_0)$  will change  $f(x)$  dramatically ( $\sim 4.5 \times$ ). However, MLP equipped with layernorm exhibits good local property.
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 \times$ ) 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 tends 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  $\mathbf{A}$  by reducing the value of improperly large entries since matrix inequality tells us  $\lambda_{\mathrm{max}}(\mathbf{A}) \leq \| \mathbf{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.
Figure 6: The effect of various regularizations and normalizations on SEEM and Q-value. LayerNorm yields a low SEEM and achieves stable Q-learning. Given the substantial disparity in the y-axis range among various regularizations, we present the results using two separate figures.
To validate the effectiveness of LayerNorm in practical offline RL settings, we select the walker2d-medium-expert-v2 task in D4RL to showcase how Q-value and SEEM evolves as the training proceeds. For comparative analysis, we also consider four popular regularization and normalization techniques, including BatchNorm [20], WeightNorm [44], dropout [45], and weight decay [35]. We
Table 1: Averaged normalized scores of the last ten consecutive checkpoints on Antmaze tasks over 10 seeds. diff-QL [51] reports the best score during the whole training process in its original paper. We rerun its official code for the average score. We also compare our method with diff-QL by the best score metrics in the parenthesis. The best score of diff-QL is directly quote from its paper.
Dataset	TD3+BC	IQL	MSG	sfBC	diff-QL	ours
antmaze-umaze-v0	40.2	87.5	98.6	93.3	95.6 (96.0)	94.3 ± 0.5 (97.0)
antmaze-umaze-diverse-v0	58.0	62.2	76.7	86.7	69.5 (84.0)	88.5 ± 6.1 (95.0)
antmaze-medium-play-v0	0.2	71.2	83.0	88.3	0.0 (79.8)	85.6 ± 1.7 (92.0)
antmaze-medium-diverse-v0	0.0	70.0	83.0	90.0	6.4 (82.0)	83.9 ± 1.6 (90.7)
antmaze-large-play-v0	0.0	39.6	46.8	63.3	1.6 (49.0)	65.4 ± 8.6 (74.0)
antmaze-large-diverse-v0	0.0	47.5	58.2	41.7	4.4 (61.7)	67.1 ± 1.8 (75.7)
average	16.4	63.0	74.4	77.2	29.6 (75.4)	80.8 (87.4)
use TD3 as the baseline algorithms. To simplify our analysis, the target network is chosen to be the current Q-network with gradient stop instead of an Exponential Moving Average (EMA) of past networks. In our experiments, it is observed that LayerNorm and WeightNorm constrain SEEM and restrain divergence. However, weight decay and dropout did not yield similar results. As illustrated in Figure 6, weight decay and dropout evolve similar to the unregularized TD3 w.o. EMA baseline. 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  $\mathrm{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 demonstrate 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,
Figure 7: The detrimental bias in policy constraint. Although strong policy constraints succeed in depressing divergence, constraining the policy around sub-optimal behavior sacrifices performance. With LayerNorm to regularize SEEM without introducing a too strong bias, our method can achieve both stable Q-value convergence and excellent performance. The shallow curves represent the trajectory of individual seeds, while the darker curve denotes the aggregate across 3 seeds.
resulting in an excellent performance. The ablation study regarding the specific configuration of adding LayerNorm can be found at Appendix F.4. The effect of other regularizations in Antmaze can be found at F.3. We also experimented with the regularizer proposed in DR3 [29]. Despite conducting 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 transition-based datasets by randomly sampling varying proportions  $(X\%)$  from the D4RL Mujoco Locomotion datasets. Here, we set several levels for  $X\in \{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
Figure 8: The performance of offline RL algorithms with the varying  $X\%$  Mujoco Locomotion dataset. As the number of transitions decreases, all algorithms have a dramatic drop in performance.
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 non-expansive, 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.
Figure 9: The performance difference between baseline with LayerNorm and without it using the same  $X\%$  dataset. The error bar represents the standard deviation over 10 seeds.
5.3 Online RL Experiments - LayerNorm Allows Online Methods without EMA.
It is natural to ask whether our analysis and solution are also applicable to online settings. Previous works [34; 13; 17; 61] has empirically 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.
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.
This work is supported by the National Key R&D Program of China (2022ZD0114900).
[1] Joshua Achiam, Ethan Knight, and Pieter Abbeel. Towards characterizing divergence in deep q-learning. arXiv preprint arXiv:1903.08894, 2019. 14
[2] Gaon An, Seungyong Moon, Jang-Hyun Kim, and Hyun Oh Song. Uncertainty-based offline reinforcement learning with diversified q-ensemble. NIPS, 2021. 2, 8
[3] Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E Hinton. Layer normalization. arXiv preprint arXiv:1607.06450, 2016. 7
[4] Leemon Baird. Residual algorithms: Reinforcement learning with function approximation. In Machine Learning Proceedings. 1995. 1, 2, 14
[5] Philip J Ball, Laura Smith, Ilya Kostrikov, and Sergey Levine. Efficient online reinforcement learning with offline data. In ICML, 2023. 7, 14
[6] Aditya Bhatt, Max Argus, Artemij Amiranashvili, and Thomas Brox. Crossnorm: Normalization for off-policy td reinforcement learning. arXiv preprint arXiv:1902.05605, 2019. 8, 14
[7] David Brandfonbrener, Will Whitney, Rajesh Ranganath, and Joan Bruna. Offline rl without off-policy evaluation. NIPS, 2021. 9, 14
[8] Huayu Chen, Cheng Lu, Chengyang Ying, Hang Su, and Jun Zhu. Offline reinforcement learning via high-fidelity generative behavior modeling. In ICLR, 2023. 2, 8, 14
[9] Lili Chen, Kevin Lu, Aravind Rajeswaran, Kimin Lee, Aditya Grover, Misha Laskin, Pieter Abbeel, Aravind Srinivas, and Igor Mordatch. Decision transformer: Reinforcement learning via sequence modeling. NIPS, 2021. 14
[10] Justin Fu, Aviral Kumar, Ofir Nachum, George Tucker, and Sergey Levine. D4rl: Datasets for deep data-driven reinforcement learning, 2020. 4, 9
[11] Justin Fu, Aviral Kumar, Matthew Soh, and Sergey Levine. Diagnosing bottlenecks in deep q-learning algorithms. In ICML, 2019. 14
[12] Scott Fujimoto and Shixiang Shane Gu. A minimalist approach to offline reinforcement learning. NIPS, 2021. 2, 8, 9, 14
[13] Scott Fujimoto, Herke Hoof, and David Meger. Addressing function approximation error in actor-critic methods. In ICML, 2018. 1, 10, 14
[14] Scott Fujimoto, David Meger, and Doina Precup. Off-policy deep reinforcement learning without exploration. In ICML, 2019. 1, 2, 14
[15] Kamyar Ghasemipour, Shixiang Shane Gu, and Ofir Nachum. Why so pessimistic? estimating uncertainties for offline rl through ensembles, and why their independence matters. NIPS, 2022. 2, 8
[16] Caglar Gulcehre, Ziyu Wang, Alexander Novikov, Thomas Paine, Sergio Gómez, Konrad Zolna, Rishabh Agarwal, Josh S Merel, Daniel J Mankowitz, Cosmin Paduraru, et al. Rl unplugged: A suite of benchmarks for offline reinforcement learning. NIPS, 2020. 9
[17] Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with a stochastic actor. In ICML, 2018. 1, 10
[18] Yizeng Han, Yifan Pu, Zihang Lai, Chaofei Wang, Shiji Song, Junfeng Cao, Wenhui Huang, Chao Deng, and Gao Huang. Learning to weight samples for dynamic early-exiting networks. In ECCV, 2022. 22
[19] Hado Hasselt. Double q-learning. NIPS, 2010. 1, 14, 23
[20] Sergey Ioffe and Christian Szegedy. Batch normalization: Accelerating deep network training by reducing internal covariate shift. In ICML, 2015. 7, 14