FANS: A Flatness-Aware Network Structure for Generalization in Offline Reinforcement Learning

Thank you for your insightful comments. To comprehensively address your concerns, we will explain the following aspects in order: inappropriate baselines, methods to reduce Q-value overestimation, ablation studies, discussions on other paradigms, and additional experiments. The order may differ from your questions, but we aim to thoroughly resolve your main concerns. We sincerely appreciate the time and effort you have invested during the rebuttal process.

> 1. Response to Inappropriate Baseline:

As pointed out in the paper Stabilizing Off-Policy Q-Learning via Bootstrapping Error Reduction, online algorithms such as TD3 and SAC often struggle to perform effectively in offline settings. This has led to the development of offline variants like TD3+BC and SAC-N. However, these methods are fundamentally built upon online algorithms, enhanced through additional regularization or structural modifications. Similarly, our proposed FANS framework is a structured enhancement tailored for offline RL. TD3+FANS and TD3+BC can thus be viewed as different extensions of the same base algorithm.

It is important to note that applying FANS on top of existing offline algorithms like TD3+BC would introduce conceptual conflicts. Specifically, BC regularization enforces conservative behavior by constraining the learned policy to stay close to the behavior policy, whereas FANS aims to improve generalization beyond the behavior policy. These objectives are inherently at odds. Therefore, we chose to apply FANS directly to the base online algorithm TD3 to better isolate and evaluate the contribution of our structured framework, without interference from other offline-specific techniques.

That said, to further address your concern, we have added experiments using TD3+BC as a baseline. The below results support our analysis and demonstrate the validity of our design choice.

Moreover, we would like to highlight that our simple combination of TD3 and FANS achieves new SOTA performance compared to a wide range of baselines, including recent SOTA methods and those focusing on generalization. Following the suggestion from Reviewer RRU4 (#R1), we have also added comparisons with several of these latest methods. Please refer to our response to #R1, point 3, for more details.

Regarding the concern about the "inappropriate baseline" in Figure 3. The primary purpose of this experiment is to provide an intuitive illustration of how FANS can significantly improve the generalization ability of a base algorithm (e.g., TD3) under limited data conditions. Figure 3 shows that while TD3 almost completely fails in this setting, FANS still achieves stable and effective performance, highlighting the potential of our structured enhancement in low-data regimes. We agree that comparing FANS with other TD3-based offline algorithms under limited data could provide a more direct evaluation of data efficiency. While we emphasize that this experiment is not intended as a comprehensive SOTA comparison, we will consider adding such comparisons in future versions to further demonstrate FANS’s effectiveness in small-data scenarios. That said, the focus of this experiment is to serve as a qualitative analysis that emphasizes the stability and generalization benefits FANS brings to the base policy. In other parts of the paper, we have already conducted systematic comparisons between FANS and various mainstream generalization methods.

> 2. Response to Reduce Q-value Overestimation:

The purpose of Figure 2 is to demonstrate the effect of the FANS framework in reducing Q-value overestimation. One common misconception is to assume that lower Q-values for OOD data are always better. In reality, what matters is the relative magnitude of Q-values $-$ i.e., whether high Q-values for OOD data are properly suppressed. Therefore, we should focus on whether the Q-value curve for OOD data exhibits signs of Q-value explosion, which is a key indicator of overestimation. In Figure 2, it is evident that the online TD3 algorithm suffers from a Q-value explosion for OOD data, while introducing FANS to TD3 effectively prevents this issue. This highlights FANS’s ability to stabilize value estimation under distribution shift.

In addition, directly comparing the absolute scale of Q-values is not always meaningful, as the values can differ by orders of magnitude due to implementation-specific factors. In FANS, the use of residual connections helps retain the original input’s scale, and the second linear layer adds new feature representations, which tends to maintain or even amplify the output magnitude. Moreover, LayerNorm normalizes each sample to have zero mean and unit variance within a feature dimension. While it does not change the overall expressive capacity, it can compress intermediate feature scales, thereby indirectly affecting the final Q-value range.

Another common misconception is that lower Q-values for OOD data necessarily lead to better performance. As shown in the MCQ paper, overly conservative handling of OOD data can suppress potential generalization, leading to suboptimal outcomes. Thus, appropriately controlling Q-values for OOD data may be more promising, although the optimal strategy remains an open question.

In summary, the goal of our analysis is to show that FANS effectively reduces harmful Q-value overestimation for OOD samples. For this reason, we directly compare TD3 and TD3+FANS in Figure 2. The online SAC curve, trained for 3 million steps, serves as a proxy for the "true" Q-value on OOD data. In contrast, directly comparing Q-values with other offline SOTA methods is not a suitable or reliable benchmark at this stage. This experiment is also designed with reference to Sec. 5.2 "Value Estimation Error" from the paper Improving Generalization in Offline Reinforcement Learning via Adversarial Data Splitting.

> 3. Response to Ablation Studies:

Our FANS framework is inspired by the well-established structured approach SimBa from the online setting. The structural differences between FANS and SimBa are as follows:

Compared to SimBa, FANS is specifically designed to enhance generalization in offline settings by additionally incorporating Gaussian activation functions and ensemble techniques. Therefore, on one hand, given that SimBa already provides a mature structured design for RL, there is no need to reiterate the analysis of the combination and effects of Residual Blocks and Layer Normalization.

We further supplement our work by presenting results of directly applying SimBa in the offline setting, i.e., TD3+SimBa, as shown in the table below:

The additional components introduced on top of SimBa $-$ the Gaussian activation function and ensemble techniques $-$ are well demonstrated in the ablation experiments in Table 4 of this paper. For instance, in the ho-me task, while the full FANS achieves 113.3, removing any single component causes severe performance drops (best ablated score: 41.3 without LayerNorm), revealing nonlinear synergies between components. Moreover, different tasks rely on distinct components (e.g., Gaussian noise is crucial for ho-m but LayerNorm dominates wa-m), proving our components dynamically collaborate rather than providing fixed incremental gains. This emergent synergy explains why the complete FANS consistently outperforms all partial implementations across diverse tasks.

Moreover, while SimBa applies residual blocks to both the actor and the critic, our FANS framework further analyzes the effects of applying residual blocks separately to the actor and the critic in offline environments, as detailed in Table 5 of this paper.

> 4. Response to Discussions on other Paradigms:

This work focuses on offline RL and does not specifically discuss other RL paradigms. However, we believe that some of the design principles embodied by FANS $-$ such as guiding the optimization process toward flatter regions through structured design and reducing estimation variance to enhance generalization $-$ may offer valuable insights for other paradigms that are also heavily affected by distribution shifts. In offline-to-online RL, continual RL, or multitask RL, a core challenge is improving the policy’s adaptability to dynamically changing environments or tasks. The modular structure of FANS provides a potential architectural framework for further exploration of these challenges.

It is important to emphasize that applying FANS to these paradigms would require appropriate extensions or adaptations tailored to their specific problem characteristics, such as incorporating task-specific memory mechanisms, adaptive modules, or dynamic sharing structures, rather than directly applying the existing architecture. Therefore, systematic investigations of FANS in other paradigms are left for future work.

> 5. Response to Additional Experiments:

We have supplemented the results on AntMaze and Adroit; please refer to our response to Reviewer RRU4 (#R1), point 4.

Dear Reviewer 6MWm:

We sincerely appreciate your insightful feedback and recognition of FANS's contribution potential (raised score to 3). We now address the key concerns you mentioned with theoretical analysis and empirical evidence:

1. Resolving Ambiguity: FANS vs. Offline Baselines in Q-value Control

New Experiment: We quantitatively measure value estimation accuracy by comparing the discounted Monte Carlo return of each algorithm’s policy trajectories (ground truth) versus its Q-value predictions. To facilitate comparison, we quantified the level of Q-value over/underestimation for each algorithm by subtracting the ground truth from the predicted Q-value and dividing the result by the ground truth. Greater/less than 0 means overestimation/underestimation. Average results across six D4RL datasets (walker2d-medium, walker2d-medium-replay, hopper-medium, hopper-medium-replay, halfcheetah-medium, halfcheetah-medium-replay) in the following table show:

As the training goes, FANS achieves the most accurate Q-value estimation. TD3BC also shows relatively accurate estimation. This proves FANS and TD3BC achieves superior value calibration – neither severe overestimating like vanilla TD3 nor severe underestimating like CQL. Crucially, while TD3BC also shows reasonable accuracy, its inferior empirical performance indicates that value accuracy alone is insufficient.

2. Revealing FANS's Fundamental Advantage: Generalization Control

To investigate why FANS outperforms TD3BC despite similar value accuracy, we analyze generalization patterns using Neural Tangent Kernel (NTK) [1]. When updating $Q$-values for $(s,a)$ using TD learning, the change at any $(\bar{s},\bar{a})$ is governed by the NTK $k_\phi(\bar{s},\bar{a},s,a)$. A simple derivation in [1] is given:

1. Parameter update for $(s,a)$:

   $\phi' = \phi + \left( \mathcal{T} Q_\phi(s, a) - Q_\phi(s, a) \right) \nabla_\phi Q_\phi(s, a)$

   ($\mathcal{T}$: target operator, learning rate omitted)

2. Q-value change at $(\bar{s},\bar{a})$ by Taylor expansion at the pre-update parameter $\phi$:

   $Q_{\phi'}(\bar{s}, \bar{a}) = Q_\phi(\bar{s}, \bar{a}) + \nabla_\phi Q_\phi(\bar{s}, \bar{a})^\top (\phi' - \phi) + \mathcal{O}(\|\Delta\phi\|^2)$

3. Generalization via NTK by plugging the second Eq to the first one:

   $Q_{\phi'}(\bar{s}, \bar{a}) = Q_\phi(\bar{s}, \bar{a}) +[ k_\phi(\bar{s},\bar{a},s,a) ] \left( \mathcal{T} Q_\phi(s, a) - Q_\phi(s, a) \right) + \mathcal{O}(\|\Delta\phi\|^2)$

   where $k_\phi(\bar{s},\bar{a},s,a)$ is $\nabla_\phi Q_\phi(\bar{s}, \bar{a})^\top \nabla_\phi Q_\phi(s, a)$

4. NTK interpretation:

   • High $k_\phi$ → Prominent generalization: $Q(\bar{s},\bar{a})$ changes significantly in sync with TD error updating on $Q(s,a)$
   • Low $k_\phi$ → Minimal generalization (tabular-like behavior when $k_\phi=0$)

Key Insight: The NTK $k_\phi(\bar{s},\bar{a},s,a)$ quantifies how much updating $Q(s,a)$ affects $Q(\bar{s},\bar{a})$*, acting as a generalization metric for TD learning.

Key Experiment: We track the normalized NTK improvement based on the NTK of initialized network (training_step=0) and explore three OOD query types of $(\bar{s}, \bar{a})$ during training:

• $(s,\pi(s)), s \in D$;
• $(s,\pi(s)+\epsilon),s \in D, \epsilon$ is noise;
• $(s’,a’) \in D$ which is next state and next action (method of DR3[2])

Due to the high dimensionality of $\nabla_\phi Q_\phi(s,a)$, the direct computation of $\underset{s, a \sim \mathcal{D}}{\mathbb{E}} |k_\phi(s, \pi(s), s, a)|$ is computationally prohibitive, a method is adopted by approximating $\Delta(\phi)$ with the contribution solely from the last layer parameters and obtain the $\underset{s, a \sim \mathcal{D}}{\mathbb{E}} |\Phi(s, \pi(s))^{\top} \Phi(s, a)|$ where $\Phi(s, a)$ signifies the representation of state-action pairs, which is the output of penultimate layer of Q-network.

Results:

We present the average results across six D4RL datasets (walker2d-medium, walker2d-medium-replay, hopper-medium, hopper-medium-replay, halfcheetah-medium, halfcheetah-medium-replay) logged every 50k step during 500K training steps (makes sure convergence) in the following tables.

Interpretation:

• The significantly lower NTK values in FANS demonstrate a fundamental suppression of pathological generalization patterns. This occurs because the kernel $k_\phi(\bar{s},\bar{a},s,a)$ acts as a generalization amplifier: High values force TD-errors from ID data $(s,a)$ to distort OOD values $Q(\bar{s},\bar{a})$. During offline RL, this propagates and amplifies extrapolation errors (especially dangerous for actions $\pi(s)$ near distribution boundaries).
• Our design directly counteracts this through: (1) residual connections that maintain stable feature baselines via identity mappings, preventing chaotic error propagation; (2) layer normalization that constrain feature co-adaptation by enforcing per-sample feature scale invariance and reducing sensitivity to aberrant activation patterns. These mechanisms collectively enhance the model’s expressive capacity and mitigates harmful generalization errors, thus shows superior performance than baselines.

3. Summary

We agree a universal criterion for "optimal" OOD suppression remains open. However, our work reveals a paradigm shift: Offline RL's extrapolation error may stem more from MLPs' poor approximation properties than algorithmic limitations. FANS addresses via two measurable advantages:

1. Appropriate Overestimation Mitigation: Prevents Q-explosion as in Fig.2
2. Generalization Safety: Minimizes NTK-based error propagation (shown above)

Thus, distinguished from existing methods that impose explicit constraints on surface-level patterns, FANS fundamentally mitigates distribution shifts by architecturally suppressing OOD overestimation and spurious generalization—constituting our primary contribution.

We provide both theoretical grounding (NTK dynamics) and multi-faceted empirical evidence (value accuracy + generalization metrics) to resolve concerns about FANS's advantages. We sincerely hope this comprehensive response clarifies FANS's unique value. Given the novel evidence and analysis, we kindly request reconsidering the score to reflect these contributions. Thank you!

[1] Reining Generalization in Offline Reinforcement Learning via Representation Distinction. NeurIPS 2023.

[2] DR3: Value-Based Deep Reinforcement Learning Requires Explicit Regularization. ICLR 2022.

Thank you for your thoughtful and inspiring comments, which will help us greatly in improving the quality of our work. We provide discussions and explanations about your concerns as follows.

> 1. Response to Motivation and the Role of Architecture:

We thank you for your insightful question. While regularizer-based methods (e.g., conservative Q-learning, behavior cloning constraints) have proven effective in offline RL by mitigating extrapolation errors and discouraging OOD actions, they primarily operate at the objective level by modifying the loss function. These methods aim to influence learning behavior via explicit penalties or constraints, often without considering the inductive biases introduced by the network architecture itself.

However, a fundamental limitation of regularizer-based approaches is their inability to directly control the geometry of the loss landscape or the sharpness of the learned solution. In the offline RL setting, where learning relies heavily on bootstrapped targets and limited data coverage, sharp or brittle solutions can lead to unstable training and poor generalization. Regularizers may reduce certain types of estimation error but do not inherently promote flatter minima or smoother optimization dynamics.

Our motivation for exploring an architectural remedy arises from this observation. The FANS framework is designed to structurally guide the optimization process toward flatter and more robust regions of the solution space. Components such as residual blocks, Gaussian activations, and layer normalization are known to promote gradient smoothness and stability, while lightweight ensemble modeling helps reduce estimation variance. These architectural choices introduce inductive biases that complement, rather than replace, regularization.

Importantly, our goal is not to position architecture as a replacement for regularization-based methods, but rather as an orthogonal and complementary perspective. In principle, architectural and objective-level approaches can be combined. In this work, however, we isolate the architectural effects to highlight their standalone benefits and to motivate further exploration of structure-aware design in offline RL.

We will clarify this motivation more explicitly in the revised manuscript.

> 2. Response to Quantifying Flatness and the Conceptual Framing:

We sincerely appreciate your insightful comments, which indeed highlight aspects not covered in the initially submitted manuscript. Thank you for bringing these to our attention. One of the core motivations of this work is to guide the optimization process toward flatter minima through structural design, thereby improving generalization in offline RL, especially under distributional shift. However, as the reviewer correctly noted, the current version of the paper does not explicitly include quantitative measures of flatness, such as the Hessian trace, eigenvalues, or sharpness, which is a limitation of our current work.

We chose not to include such analysis in the main paper for the following reasons: First, in RL, especially when using bootstrapped targets like Q-value estimates, the loss surface is highly non-stationary, making the direct computation of the Hessian unstable and computationally expensive. Second, this work focuses on understanding the role of architecture itself by isolating it from objective-level interventions. Therefore, we primarily evaluated improvements in OOD generalization, estimation variance, and overall policy performance as indirect evidence of improved flatness and robustness.

We fully agree that the causal link between architectural design, loss surface flatness, and policy robustness should be more rigorously demonstrated. Thus, in the revised version, we will supplement the paper with additional analyses to characterize the flatness of the loss landscape. Since we cannot include additional figures or appendices in the rebuttal, we have instead described our methodology for measuring landscape flatness in Reviewer jCXx (#R2)’s comment 3 for reference.

We extend our sincere gratitude for your inspiring and professional insights, which are invaluable in enhancing the quality of our work. In response to your concerns, we report below the explanations and answer your questions accordingly.

> 1. Response to Computational Overhead of FANS:

We provided a detailed analysis of FANS’s computational overhead in our response to Reviewer RRU4 (#R1), point 2. We hope this is helpful to you.

> 2. Response to Hyperparameter Count and Sensitivity:

FANS introduces minimal additional hyperparameters compared to standard offline RL baselines. Its architectural components are largely self-contained and do not require extensive tuning. The table below summarizes the key components introduced by FANS and whether they introduce new tunable hyperparameters.

FANS also exhibits low sensitivity to hyperparameter choices. The use of residual connections and layer normalization enhances training stability and enables deeper architectures without introducing optimization instability. The fixed Gaussian activation promotes smoother gradients and helps reduce sensitivity to learning rates and batch sizes. Overall, FANS maintains robust performance across a broad range of hyperparameter settings, alleviating the need for exhaustive tuning.

> 3. Response to Flatter Regions of the Loss Landscape:

Your suggestion is highly valuable. Considering that we are unable to submit a PDF response, we outline below a proposed approach to illustrate whether FANS effectively guides training dynamics toward flatter regions of the loss landscape:

(1) Selection of Converged Policies:
We identify a point in training where the loss stabilizes and performance metrics (e.g., policy returns) no longer change significantly, indicating convergence. At this stage, we fix the parameters $\phi$ of the policy network $\pi_\phi$.

(2) Local Surface Construction:
We randomly sample a batch of states from the replay buffer and apply small perturbations to each state. For each perturbed state, we compute the corresponding action using the fixed policy network $\pi_\phi$. The original (unperturbed) state's action is treated as the reference optimal action.

(3) Surface Visualization:
We construct a 2D surface where the x- and y-axes correspond to the direction and magnitude of perturbations in the state space, and the z-axis represents the normalized Euclidean distance between the perturbed action and the reference action. Concave points on this surface correspond to the optimal actions. A flatter surface indicates that the policy output is less sensitive to input perturbations, suggesting greater local robustness and a smoother loss landscape.

(4) Comparative Analysis:
We perform this analysis on policies trained with and without FANS. By comparing the resulting action surfaces, we aim to visually demonstrate that FANS leads to a flatter, smoother local landscape $-$ supporting the claim that it enhances training stability and generalization.

Although directly visualizing the high-dimensional loss landscape in reinforcement learning is inherently difficult, this approach serves as an interpretable proxy. We will include such comparative plots in the final version of the paper to strengthen the empirical support for FANS.

> 4. Response to Weakness: Inappropriate Baseline in Figure 3:

Thank you for your valuable feedback. Reviewer 6MWm (#R4) has raised a similar concern, and we kindly refer you to our detailed response to #R4, point 1, for your reference.

Thanks for your supportive feedback! We will include the supplymentary results in our revised paper according to your and other reviewers' suggestions to further enhance the completeness of the work.

We sincerely appreciate your thoughtful and professional insights, which have greatly improved the quality of our work. In response to your concerns, we provide the explanations and answers below.

> 1. Response to Lack of Ablation Studies:

We have conducted ablation studies in Section 4.4 that systematically evaluate the contribution of each module within FANS to the overall performance. Additionally, Appendix B.2 provides further structural analysis, verifying the effectiveness of each component through parameter variations. For a more detailed explanation of the ablation experiments, please also refer to our response to Reviewer 6MWm (#R4), point 3. We hope these clarifications help to better illustrate the relative importance of the components.

> 2. Response to the Time Complexity and Computational Overhead of FANS:

(1) Residual Block Complexity:

Each residual block in FANS (as shown in Table 1 of the main paper) comprises two linear layers, a LayerNorm, a Gaussian activation, and a residual connection.

• The dominant operations are the two linear layers, each with a time complexity of $\mathcal{O}(d^2)$, where $d$ is the hidden dimension.
• LayerNorm, the Gaussian activation ($\exp(-x^2)$), and the residual addition are all element-wise operations with complexity $\mathcal{O}(d)$.

Thus, each block maintains an overall complexity of $\mathcal{O}(d^2)$, consistent with standard MLP layers used in typical offline RL networks. The additional operations increase the constant factor but do not alter the asymptotic complexity. Additionally, these Residual Blocks are only applied to the critic network, remaining lightweight and efficient.

(2) Ensemble Modeling:

We acknowledge that ensemble modeling introduces additional cost, scaling linearly with the number of models $M$, resulting in $\mathcal{O}(M \cdot d^2)$ complexity. However:

• The ensemble is only used in the critic, not the actor;
• We adopt a small $M$ (e.g., 2-5) while still achieving strong performance $-$ significantly fewer than SAC-N, which requires 100-500 Q-networks, and EDAC, which requires 20-50 Q-networks;

(3) Practical Impact:

Despite the slight increase in per-iteration cost, FANS demonstrates stronger generalization and stability, requiring fewer training epochs overall. In Figure 6 of the Appendix, we show that FANS converges faster than strong baselines under the same compute budget.

Summary:

In summary, while FANS introduces moderate constant-factor overhead due to its structured architectural design, the overall time complexity remains $\mathcal{O}(d^2)$, matching conventional offline RL networks. The improvements in generalization and training robustness justify this trade-off. Bellow the table "Time Complexity and GPU Consumption Comparison":

We will include detailed time complexity analysis and empirical runtime comparisons in the revised version.

> 3. Response to Missing Experimental Comparison with state-of-the-art (SOTA) Methods:

Thank you for the valuable suggestion. As for A2PR, we directly extracted their MuJoCo Gym results from their original paper. The comparison between FANS and A2PR is summarized in the table below, where FANS outperforms A2PR on 7 out of 9 tasks. We will include this comparison in the revised manuscript’s main text.

Regarding FQL and SORL, we note that neither method reports results on the standard MuJoCo benchmark, which serves as the primary evaluation suite in our study. As a result, a direct empirical comparison was not feasible under our current experimental setup. Nevertheless, we greatly value these contributions and will explicitly include a discussion of both methods in the Related Work section of the revised version, clarifying how they relate to our approach. We are grateful to you for pointing out these important works.

> 4. Response to Insufficient Main Experimental Results:

Due to the limited time available during the rebuttal phase, we have conducted a preliminary set of experiments on the AntMaze and Adroit tasks. The results are shown in the table below. More comprehensive results will be provided in the revised version of the paper.

(1) AntMaze:

(2) Adroit: