Adaptive Neighborhood-Constrained Q Learning for Offline Reinforcement Learning

We appreciate the time and effort you have dedicated to providing feedback on our paper and are grateful for the meaningful comments.

Q1: The bound in Proposition 1 depends on $|\mathcal{S}|$.

Thank you for the valuable comment. We first clarify the original bound and then provide a tighter version that eliminates the $|\mathcal{S}|$-dependency.

(1) Original bound: The bound in Proposition 1 is mathematically valid, but potentially loose. It scales linearly with the neighborhood radius $\epsilon$, which makes it meaningful when $\epsilon$ is small. However, the dependence on $|\mathcal{S}|$ can indeed make the bound less practical in large-scale settings. This linear dependence on $|\mathcal{S}|$ originates from the pointwise Lipschitz assumption on the transition dynamics, i.e.,

$|P(s'|s,a_1) - P(s'|s,a_2)| \leq K_P \\|a_1 - a_2\\|, \quad \forall s,s' \in \mathcal{S},\forall a_1,a_2 \in \mathcal{A}$

In the proof (Eq. (35)), this pointwise difference is summed over all possible next states, thus introducing the factor $|\mathcal{S}|$.

(2) Tighter bound without $|\mathcal{S}|$-dependency:

The $|\mathcal{S}|$-term can be effectively eliminated by replacing the pointwise Lipschitz assumption with an $L_1$-Lipschitz assumption on the transition dynamics:

$\\|P(\cdot|s, a_1)-P(\cdot|s, a_2)\\| _ 1 \leq K_{L1} \\|a_1 - a_2\\|, \quad \forall s \in \mathcal{S},\forall a_1,a_2 \in \mathcal{A}$

This assumption is consistent with prior theoretical RL studies [5], where bounding distributional shift under action perturbations is critical for convergence guarantees. A small $K_{L1}$ is also empirically supported by observations in physical systems, where small action changes typically lead to limited next-state variation [6].

In the proof, we replace Eq. (35) with:

$\sum_s \left| P(s \mid s', \pi_1(s')) - P(s \mid s', \pi_2(s')) \right| \leq K_{L1} \\|\pi_1(s') - \pi_2(s')\\|.$

The rest of the derivation remains unchanged, leading to a revised bound:

$\mathrm{D_{TV}}\left(d^{\pi_1}(\cdot), d^{\pi_2}(\cdot)\right) \leq \frac{\gamma K_{L1}}{2(1 - \gamma)} \epsilon$.

This new bound is more broadly applicable for large-scale problems, as it is independent of $|\mathcal{S}|$ and typically yields smaller constants in real-world applications. We thank the reviewer for this observation and will include the new bound in the revision to further strengthen Proposition 1.

Q2: About the meaning of $S$ in Theorem 1 and fixed $\nu$.

We apologize for the confusion caused by the insufficient explanation of notations. Since both the support and neighborhood constraint set are defined over actions conditioned on a given state, Theorem 1 analyzes their approximation relationship at a fixed state, focusing on their difference in the action space. We adopt following abstract notations to formalize the analysis.

Theorem 1. Let $S \subseteq \mathbb{R}^d$ be the compact support of a distribution $\nu$, and let $X_1,\dots,X_n$ be i.i.d. samples from $\nu$. Define $U_{n,\epsilon} = \bigcup_{i=1}^n B(X_i, \epsilon)$ as the union of closed balls of radius $\epsilon$ centered at the samples. Let $\mathcal{N}(S, \epsilon/2)$ denote the covering number of $S$…

In the context of offline RL: $\nu$ represents the behavior policy distribution on a given state, which is generally assumed fixed in offline RL. $S$ (distinct from the state space $\mathcal{S}$) is defined as the support of behavior policy $\nu$ at that state. $X_1,\dots,X_n$ are independent and identically distributed (i.i.d.) samples from the behavior policy $\nu$, corresponding to dataset actions at that state, consistent with the standard offline RL assumption of i.i.d. data from a fixed behavior policy. Thus, the theorem indeed concerns actions and $\nu$ is fixed.

Q3: Lemma 1 is a very standard result and can be seen from e.g. [Achiam et al. 2017]. Lemma 2 comes from [43].

Thank you for pointing this out. We agree that Lemma 1 and Lemma 2 are based on established results. Lemma 1 is derived from the performance difference lemma [1] and is related to results in TRPO [2] and CPO [3]. Lemma 2 is a direct corollary of Lemma 4 in [4], specialized to the case of action extrapolation, as noted in the proof. These lemmas are included not as novel contributions, but to highlight the limitations of existing approaches and to motivate the design of our method. We present them as lemmas (rather than propositions or theorems) to distinguish them from our contributions. However, we acknowledge that their connections to prior work were not sufficiently explained in the main text, which could cause confusion. We will revise the manuscript to clarify these points and properly attribute the prior results. The main theoretical contributions of our work are captured in the proposition and theorem.

Q4: The first half of the technical sections (before Definition 5) and the second seem to be detached.

Thanks for the comment. The flow of the technical sections is as follows:

(1) The first half (before Definition 5) analyzes the limitations of existing constraint types, proposes neighborhood constraint as a new type, and establishes its theoretical properties under uniform neighborhood radii.

(2) The second half (after Definition 5) develops a tractable optimization algorithm that enables Q learning under the proposed neighborhood constraint. Notably, although the optimization algorithm is presented specifically for the adaptive neighborhood in Definition 5, i.e., using radius $\epsilon \exp(-\alpha A(s,a))$, the framework is general and can accommodate any adaptive radius schemes. Specifically, one can simply replace $\epsilon \exp(-\alpha A(s,a))$ in Definition 5 (and Eq. (11)) with $\epsilon f(s,a)$ to define arbitrary per-sample neighborhood radii, where $f: \mathcal S \times \mathcal A \to \mathbb R^+$ is any function adjusting the radius. Correspondingly, in Eq. (12), $\exp(\alpha A(s,a))$ becomes $1/f(s,a)$, and Eq. (13) becomes:

$\max_{\mu_\omega} \mathbb E_{(s,a)\sim \mathcal{D}} [Q_\theta(s,a+\mu_\omega(s,a)) - \lambda \\|\mu_\omega(s,a)\\| / f(s,a)].$

All other equations remain unchanged, yielding an algorithm that supports arbitrary neighborhood radius schemes.

(3) Definition 5 serves as a transition point between the theoretical and practical parts. It provides a concrete instantiation of adaptive neighborhoods, where the radius is adjusted based on estimated advantage values, motivated by empirical insights.

Q5: About Definition 5 and advantage estimates.

Definition 5 does not seem to be very well coupled with previous developments: the authors did not sufficiently motivate the radius $\epsilon \exp(-\alpha A(s,a))$ nor explained how it connects to the covering number theorem.

(1) Definition 5 is introduced as a practical instantiation of adaptive neighborhoods, where the sample-wise radius is set as $\epsilon \exp(-\alpha A(s,a))$ based on estimated advantages. As explained in the paragraph preceding Definition 5, this design is motivated by empirical insights: expert-like actions are typically concentrated and require smaller neighborhood to reduce extrapolation, while suboptimal actions tend to be more dispersed and can tolerate and benefit from larger neighborhood to enable broader search. The exponential form of advantages $\epsilon \exp(-\alpha A(s,a))$ heuristically captures this intuition.

(2) Definition 5 is intended as a practical complement based on the theoretical foundation established under the uniform setting. Theorem 1 establishes theoretical guarantees under uniform neighborhood radii, serving as a foundation for understanding geometric coverage. Building on this foundation, Definition 5 extends the formulation by introducing an adaptive radius scheme guided by empirical insights. It serves as an empirically grounded instantiation that fits naturally within our framework and supports the broader goal of enabling flexible, pointwise conservatism.

More importantly, this quantity ($\epsilon \exp(-\alpha A(s,a))$) relies on accurate estimates of advantage in the offline context. But if we can have such accurate estimates free of OOD issues, in principle we don't have to employ this perturbation scheme?

Thanks for the comment. We will clarify from two perspectives:

(1) In-distribution estimation: This quantity (adaptive radius) relies solely on advantage estimates at dataset points $(s,a) \in \mathcal D$, which are in-distribution (in fact, in-sample) and therefore more reliable than estimates at OOD points.

(2) No need for precise advantage: The purpose of using advantage is to qualitatively distinguish between good and suboptimal actions, and exact advantage values are not required. The exponential form $\epsilon \exp(-\alpha A(s,a))$ is a soft heuristic to bias the neighborhood radius. Our motivation does not assume or require the radius to be exactly proportional to the exponential of the advantage.

Reference

[1] Kakade et al. Approximately optimal approximate reinforcement learning. ICML 2002.

[2] Schulman et al. Trust region policy optimization. ICML 2015.

[3] Achiam et al. Constrained policy optimization. ICML 2017.

[4] Li et al. When Data Geometry Meets Deep Function: Generalizing Offline Reinforcement Learning. ICLR 2023.

[5] Asadi et al. Lipschitz continuity in model-based reinforcement learning. ICML 2018.

[6] Todorov et al. Mujoco: A physics engine for model-based control. IROS 2012.

We appreciate the time and effort you have dedicated to providing feedback on our paper and are grateful for the meaningful comments.

Q1: A more detailed analysis of hyperparameter sensitivity would be beneficial to understand the robustness of the method.

Thank you for the suggestion. We analyze the hyperparameter sensitivity from the following three aspects:

(1) Superior performance within moderate parameter ranges. As shown in the ablation study in Section 4.4, ANQ generally exhibits stronger and more robust performance with intermediate values of Lagrange multiplier $\lambda$ and inverse temperature $\alpha$. Since $\lambda=0$ and $\lambda=\infty$ correspond to no constraint and sample constraint respectively, the results indicate that the proposed neighborhood constraint outperforms both across a range of hyperparameters. Similarly, as $\alpha=0$ corresponds to a fixed neighborhood radius, the results show that advantage-based adaptive neighborhoods generally outperform uniform neighborhoods.

(2) Relatively limited hyperparameter tuning. The specific ANQ algorithm involves relatively little hyperparameter tuning. As detailed in Appendix C1, we fix $\alpha=1$ across all tasks, which yields consistently strong performance. For $\lambda$, we use only a small set of values: $\lambda = 5$ for Antmaze, and $\lambda \in \\{0.1, 5\\}$ for Gym locomotion, involving less tuning compared to baselines like CQL and SPOT.

(3) Insights into hyperparameter selection. For selecting $\lambda$, a general principle is as follows: for narrow-distribution datasets (typically expert or demonstration data), a larger $\lambda$ can be used to induce smaller neighborhoods, thereby helping to more effectively control extrapolation error. For more dispersed datasets (often containing noisy or mixed-quality data), the impact of extrapolation error is relatively weaker; hence, a smaller $\lambda$ (i.e., larger neighborhoods) can be adopted to promote broader optimization over the action space and mitigate the impacts of suboptimal actions. Additionally, using a larger $\alpha$ in such cases can further enhance this effect by downweighting low-advantage actions more aggressively.

Q2: Some new sota support/sample constraint baselines are suggested such as [1][2].

Thanks for the suggestion. The table below reports the overall performance of STR, CPI, and ANQ. ANQ achieves the best results in the Antmaze domain and demonstrates competitive performance with CPI in the Gym locomotion domain. A full comparison will be included in the revised manuscript.

Table 2: Total scores on Gym locomotion and Antmaze tasks, averaged over 5 random seeds.

Q3: Ablation that changes $\mu_\omega$ to a random variable sampled from Gaussian.

Thanks for the insightful question.

(1) Clarification on exploiting rather than smoothing. Defining the auxiliary policy $\mu$ as Gaussian noise is reminiscent of the target policy noise trick introduced in TD3. Specifically, in TD3, Gaussian noise is added to the policy actions in the Bellman target. Because the trained policy is prone to overfitting local optima in the Q-function landscape, such noise injection can smooth the value estimates and reduce accumulated errors in Q-function training. By contrast, in our method, $\mu_\omega$ is applied to fixed in-dataset actions. Therefore, there is no need for Q-function estimate smoothing as in TD3. On the contrary, our method exploits, rather than smooths, the local Q landscape surrounding dataset actions, using it as a guide for the inner optimization. In our algorithm, $\mu_\omega$ is optimized to seek better actions within the neighborhood of dataset actions. Replacing $\mu_\omega$ with Gaussian noise essentially turns this directed search into undirected random perturbations. Thus, our method is not intended to smooth the Q landscape, but rather to exploit its local structure for improvement.

(2) Empirical results. We compared different choices of the auxiliary policy $\mu$ on the Gym locomotion tasks. Results are shown in the table below. The tested variants include:

• ANQ w/o $\mu$: $\mu$ is set to zero, equivalent to IQL.
• ANQ w/ Gaussian noise $\mu$: $\mu$ is replaced with Gaussian noise, i.e., $\mu \sim \text{clip}(\mathcal{N}(0, 0.04), -0.5, 0.5)$ as used in TD3 and TD3+BC.
• ANQ w/ optimized $\mu$: our default ANQ algorithm where $\mu_\omega$ is optimized.

Table 1: Averaged normalized scores on Gym locomotion tasks over 5 random seeds.

The results show that using an optimized $\mu$ significantly outperforms the other two variants, especially on low-quality datasets, clearly demonstrating the benefit of our design. Additionally, ANQ with Gaussian noise performs slightly worse than ANQ without $\mu$, likely because the Q values of dataset actions do not require smoothing, and the randomness introduced by Gaussian noise may cause the Bellman target to select actions inferior to the original dataset actions.

We appreciate the time and effort you have dedicated to providing feedback on our paper and are grateful for the meaningful comments.

Q1: Could the authors comment how/whether it is possible to extend to stochastic policies?

Thank you for the meaningful question.

(1) From a practical perspective, our algorithm is fully compatible with stochastic policies. Both policies introduced in the method (the auxiliary policy and the final policy) can be replaced by stochastic counterparts without modification to the core framework.

• For optimizing a stochastic auxiliary policy $\mu_\omega$, Eq. (13) can be directly optimized using the reparameterization trick:

$\max_{\mu_\omega} \mathbb{E}_ {(s,a)\sim \mathcal{D},\epsilon\sim\mathcal{N}(0,1)} \left[ Q_\theta(s,a+\mu_\omega(\epsilon;s,a)) - \lambda \exp(\alpha (Q_{\theta'}(s,a) - V_\psi(s))) \\|\mu_\omega(\epsilon;s,a)\\| \right].$

• For extracting a stochastic final policy $\pi_\phi$, the regression problem in Eq.(16) can be reformulated as a maximum likelihood problem:

$\min_{\pi_\phi} \mathbb{E}_ {(s,a)\sim \mathcal{D},\epsilon\sim\mathcal{N}(0,1)} \exp (\beta(Q_{\theta'}(s,a+\mu_\omega(\epsilon;s,a)) - V_\psi(s))) \log \pi_\phi(a+\mu_\omega(\epsilon;s,a)|s)$

(2) Empirically, we tested a stochastic version of the algorithm, denoted as ANQ-Gaussian, where both policies are Gaussian. The results, reported in the table below, show that ANQ-Gaussian performs comparably to the original ANQ across most Gym locomotion tasks, though slightly worse on a few.

Table 1: Averaged normalized scores on Gym locomotion tasks over 5 random seeds.

(3) Theoretically, while Proposition 1 relies on a deterministic policy, Lemma 2 and Theorem 1 hold independently of policy determinism, as they concern properties of the neighborhood constraint itself. However, in practice, a stochastic policy may occasionally generate actions that violate the neighborhood constraint, making it only probabilistically satisfied. Investigating the theoretical properties of probabilistic neighborhood constraints presents an interesting direction for future work.

Q2: Sensitivity to $\lambda$ and $\alpha$.

We appreciate the comment and will discuss the hyperparameter sensitivity in three aspects:

(1) Superior performance within moderate parameter ranges. As shown in the ablation study, ANQ generally exhibits stronger and more robust performance with intermediate values of Lagrange multiplier $\lambda$ and inverse temperature $\alpha$. Since $\lambda=0$ and $\lambda=\infty$ correspond to no constraint and sample constraint respectively, the results indicate that the proposed neighborhood constraint outperforms both across a range of hyperparameters. Similarly, as $\alpha=0$ corresponds to a fixed neighborhood radius, the results show that advantage-based adaptive neighborhoods generally outperform uniform neighborhoods.

(2) Relatively limited hyperparameter tuning. The specific ANQ algorithm involves relatively little hyperparameter tuning. As detailed in Appendix C1, we fix $\alpha=1$ across all tasks, which yields consistently strong performance. For $\lambda$, we use only a small set of values: $\lambda = 5$ for Antmaze, and $\lambda \in \\{0.1, 5\\}$ for Gym locomotion, involving less tuning compared to baselines like CQL and SPOT.

(3) Insights into hyperparameter selection. For selecting $\lambda$, a general principle is as follows: for narrow-distribution datasets (typically expert or demonstration data), a larger $\lambda$ can be used to induce smaller neighborhoods, thereby helping to more effectively control extrapolation error. For more dispersed datasets (often containing noisy or mixed-quality data), the impact of extrapolation error is relatively weaker; hence, a smaller $\lambda$ (i.e., larger neighborhoods) can be adopted to promote broader optimization over the action space and mitigate the impacts of suboptimal actions. Additionally, using a larger $\alpha$ in such cases can further enhance this effect by downweighting low-advantage actions more aggressively.

Please remember to respond to authors' rebuttal as soon as possible.

Thank you!

-AC

We appreciate the time and effort you have dedicated to providing feedback on our paper and are grateful for the meaningful comments.

Q1: About the cost and complexity of bi-level optimization.

(1) Runtime efficiency.

We briefly describe the training time results in the Benchmark Results section, with detailed comparisons provided in Figure 5 of Appendix D.1. The proposed ANQ is among the fastest tier of offline RL algorithms, on par with efficient baselines such as AWAC, IQL, and TD3BC. Several factors contribute to this efficiency:

• ANQ is model-free and ensemble-free, making it more efficient than model-based (e.g., MOPO [1]) and ensemble-based methods (e.g., EDAC [2]);
• ANQ uses simple three-layer MLPs for policy and value networks (as in TD3), which are more lightweight than complex architectures used in methods like Decision Transformer [3] or Diffusion Q-learning [4];
• ANQ optimizes deterministic policies with a policy update frequency of 2 (as in TD3), offering slightly better efficiency than algorithms with stochastic policies;
• ANQ solves the bi-level objective via alternating updates, requiring only one extra lightweight optimization step for the auxiliary policy compared to IQL.

As a result, ANQ achieves competitive training efficiency while maintaining strong performance.

(2) Memory usage.

All four networks used in ANQ are three-layer MLPs, resulting in low memory overhead as well. We test GPU memory usage on the halfcheetah-medium-replay-v2 dataset using a GeForce RTX 3090. As shown in the table below, the memory consumption of ANQ is also comparable with other efficient offline RL algorithms.

Table 2: GPU memory usage of offline RL algorithms.

Q2: About potential instability from bi-level optimization.

Thanks for the thoughtful comment. As shown in the training curves (Appendix D.2, Figures 6 and 7), ANQ generally demonstrates stable learning and convergence across most tasks. We observe no notable signs of instability such as divergence or diminishing returns.

this could lead to issues with convergence or training stability, particularly in environments with high variance or limited data.

(1) High-variance environments. Hopper and Walker2d are commonly used to evaluate robustness in high-variance environments due to their sensitive dynamics. Even small action changes can cause the agent to fall, leading to highly variable returns. Our results show that ANQ maintains relatively stable training progress and strong convergence performance on both tasks, suggesting that the bi-level optimization framework does not compromise stability in these high-variance environments.

(2) Limited data settings. It is worth noting that Section 4.3 specifically evaluates performance under limited data. As shown in Figure 2, ANQ consistently outperforms baselines with more stable final performance across varying data discard ratios, supporting its robustness in limited data scenarios.

A more detailed discussion of the potential instability and its mitigation would provide clearer guidance for users and researchers adopting this method.

While our results demonstrate ANQ's overall stability across tasks, we agree that guidance on mitigating potential instability can be valuable. Below are two practical suggestions based on our analysis.

(1) Increase $\lambda$: Ablation results show that too small $\lambda$ (i.e., too large neighborhoods) can lead to increased return variance. Therefore, increasing $\lambda$, which tightens the neighborhood constraint, may help enhance stability if necessary, though potentially at the cost of reduced policy flexibility.

(2) Decrease $\alpha$: Although ANQ remains stable across a wide range of $\alpha$ in our experiments, we hypothesize that setting a smaller $\alpha$ may further enhance stability by preventing low-quality data points from having excessively large neighborhood radii.

Q3: Effects and more insights of hyperparameter tuning.

Does their choice ($\lambda$ and $\alpha$) significantly impact training stability or the final policy performance?

(1) Effects of $\lambda$ and $\alpha$. We have conducted an ablation study in Section 4.4 on the Lagrange multiplier $\lambda$ and inverse temperature $\alpha$, analyzing learned Q-values and performance. ANQ generally exhibits stronger and more robust performance with intermediate values of $\lambda$ and $\alpha$. Since $\lambda=0$ and $\lambda=\infty$ correspond to no constraint and sample constraint respectively, the results indicate that the proposed neighborhood constraint outperforms both across a range of hyperparameters. Similarly, as $\alpha=0$ corresponds to a fixed neighborhood radius, the results show that advantage-based adaptive neighborhoods generally outperform uniform neighborhoods.

Could the authors provide more insight into how these hyperparameters are tuned in practice?

(2) Hyperparameter details. The ANQ algorithm involves relatively little hyperparameter tuning. As detailed in Appendix C1, we fix $\alpha=1$ across all tasks, which yields consistently strong performance. For $\lambda$, we use only a small set of values: $\lambda = 5$ for Antmaze, and $\lambda \in \\{0.1, 5\\}$ for Gym locomotion, involving less tuning compared to baselines like CQL and SPOT.

(3) Insights into hyperparameter selection. For selecting $\lambda$, a general principle is as follows: for narrow-distribution datasets (typically expert or demonstration data), a larger $\lambda$ can be used to induce smaller neighborhoods, thereby helping to more effectively control extrapolation error. For more dispersed datasets (often containing noisy or mixed-quality data), the impact of extrapolation error is relatively weaker; hence, a smaller $\lambda$ (i.e., larger neighborhoods) can be adopted to promote broader optimization over the action space and mitigate the impacts of suboptimal actions. Additionally, using a larger $\alpha$ in such cases can further enhance this effect by downweighting low-advantage actions more aggressively.

Is there an automatic way to tune these parameters based on the environment or data?

(4) Automatic tuning. As with many offline RL algorithms, these parameters primarily control the trade-off between conservatism and exploration. Automatically tuning such parameters remains a well-recognized and important open problem in offline RL. Research progress in this direction could be integrated into our method by leveraging the above hyperparameter selection insights. While automatic tuning is not the main focus of this work, we believe it is a promising direction for future research.

Q4: While the method performs well in standard benchmark tasks, how does it scale to high-dimensional action spaces or environments with partially observable states?

Thanks for the good question.

(1) Partially observable environments. To our knowledge, partially observable settings are not commonly featured in standard offline RL benchmarks. To simulate partial observability, we modify the Hopper and Walker2d environments by removing the agent's z-axis position, mimicking a height sensor failure. We compare representative algorithms with different constraint types, including TD3+BC, IQL, SPOT, and ANQ. Results averaged over 5 random seeds are shown in the table below. While all methods exhibit performance drops under partial observability, ANQ outperforms the baselines on most datasets and achieves the highest total score, demonstrating its robustness in such settings.

Table 1: Normalized scores in partially observable environments without height information, averaged over 5 random seeds.

(2) Complex tasks and high-dimensional action spaces.

AntMaze represents one of the most challenging offline RL benchmarks, with high-dimensional action space and sparse-reward, goal-reaching nature. The agent must not only coordinate multiple joints for locomotion but also compose suboptimal trajectory fragments to reach the goal. Empirical results show that ANQ significantly outperforms baseline methods on these tasks, especially in the most complex large-maze variants, demonstrating its effectiveness in such complex settings.

From a method design perspective, the proposed neighborhood constraint may be particularly beneficial in high-dimensional action spaces. In such settings, support constraint methods often struggle to model the behavior policy accurately, while sample constraint methods risk being overly conservative due to limited coverage of near-optimal actions. ANQ avoids both issues by enabling controlled generalization without behavior policy modeling.

Reference

[1] Yu et al. Mopo: Model-based offline policy optimization. NeurIPS 2020.

[2] An et al. Uncertainty-based offline reinforcement learning with diversified q-ensemble. NeurIPS 2021.

[3] Chen et al. Decision transformer: Reinforcement learning via sequence modeling. NeurIPS 2021.

[4] Wang et al. Diffusion policies as an expressive policy class for offline reinforcement learning. ICLR 2023.

Please remember to respond to authors' rebuttal as soon as possible.

Thank you!

-AC