SlowGuess/NeurIPS_2023_OCR huggingface dataset:
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.
Table 5: The effect of LayerNorm implementations on two challenging Antmaze tasks.
Dataset	w.o. LN	LN(0)	LN(0,1,)	LN(1,2)	LN(2)	LN(0,1,2)	LN(no learnable)
antmaze-large-play-v0	1.6	0	0	8.3	17.8	72.7	72.8
antmaze-large-diverse-v0	4.4	60.2	68	77.1	65.5	66.5	66.7
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.
Figure 14: Baird's Counterexample.
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} = \gamma G_{\theta_{t}}(X_{t}^{*}, X) - G_{\theta_{t}}(X, X) \approx (\gamma - 1)G_{\theta_{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_{\theta_{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} = \gamma G_{\theta_{t}}(X_{t}^{*}, X) - G_{\theta_{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), ReLU-activated MLP without any norm layer becomes a linear function for the points that are excessively out of the dataset range. This fact can be comprehended in the following way: Consider an input  $\lambda x$ , the activation state for every neuron in the network becomes deterministic when  $\lambda \rightarrow \infty$ . 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  $\lambda x_0, \lambda \geq C$ , this will cause an arbitrarily large NTK value between extreme points  $\lambda x_0$  and dataset point  $x_0$ . Because now  $f_{\theta}(\lambda x)$  can be written in equivalent form  $W^T(\lambda x)$ , so  $\phi(\lambda x) = \lambda W$  becomes linear proportional to  $\lambda$ . This further induces linear NTK value between  $x$  and  $\lambda x$  that is unbounded. This phenomenon can be intuitively interpreted as below, since the value prediction of  $f_{\theta}(x)$  on line  $\lambda x_0$  is almost a linear function  $W^T x$  for large  $\lambda$ , any subtle change of effective parameter  $W$  will induce a large change on  $\lambda x_0$  for large  $\lambda$ . This demonstrates a strong correlation between the value prediction between point  $x_0$  and far-away extreme point  $\lambda x_0$ . The rigorous proof can be found in Proposition 1.
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.
Figure 15: An observation that policy actions tend to gravitate toward extreme points when Q-value divergence occurs. In D4RL environments, where action spaces are bounded, the ratio of the policy action norm to the norm of points at the boundary (i.e., extreme points) provides an indication of whether the current action is approaching these extreme points. To facilitate the observation of Q-value divergence, the experimental setup employed is consistent with that described in Section 3 of the paper. From the figure, it can be seen that the ratios for policy actions are nearly 1, suggesting a close proximity to the extreme points. In contrast, the ratios for actions within the dataset are significantly lower than 1.
In Assumption 2, we posit that the direction of NTK and the policy remains stable following a certain period of training. We validate 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.
Figure 16: NTK similarity, action similarity, model parameter similarity, and model parameter norm curves in D4RL Mujoco Walker2d tasks.
Figure 17: NTK similarity, action similarity, model parameter similarity, and model parameter norm curves in D4RL Mujoco Halfcheetah and Hopper tasks.
Figure 18: NTK similarity, action similarity, model parameter similarity, and model parameter norm curves in Antmaze tasks.
Figure 19: The normalized kernel matrix's SEEM (in red) and the estimated Q-value (in blue) in D4RL Mujoco tasks. For each environment, results from three distinct seeds are reported.
antmaze-daze-diverse-v0 (3 seeds)
(a) The normalized kernel matrix's SEEM (in red) and the estimated Q-value (in blue) in D4RL Mujoco tasks.
(b) The unnormalized kernel matrix's SEEM. The three curves in each environment correspond directly to those presented in Figure (a)
Figure 20: In Figure (a), an inflation in the estimated Q-value coincides with 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.