Offline Reinforcement Learning with OOD State Correction and OOD Action Suppression

Authors: yixiu  mao
qi  wang
chen  chen
yun  qu
xiangyang  ji
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-05 14:49:08

Abstract: In offline reinforcement learning (RL), addressing the out-of-distribution (OOD) action issue has been a focus, but we argue that there exists an OOD state issue that also impairs performance yet has been underexplored. Such an issue describes the scenario when the agent encounters states out of the offline dataset during the test phase, leading to uncontrolled behavior and performance degradation. To this end, we propose SCAS, a simple yet effective approach that unifies OOD state correction and OOD action suppression in offline RL. Technically, SCAS achieves value-aware OOD state correction, capable of correcting the agent from OOD states to high-value in-distribution states. Theoretical and empirical results show that SCAS also exhibits the effect of suppressing OOD actions. On standard offline RL benchmarks, SCAS achieves excellent performance without additional hyperparameter tuning. Moreover, benefiting from its OOD state correction feature, SCAS demonstrates enhanced robustness against environmental perturbations.

Main paper: 
 introduction: 
 Deep reinforcement learning (RL) shows promise in solving sequential decision-making problems, gaining increasing interest for real-world applications [42, 57, 63, 53, 7] . However, deploying RL algorithms in extensive scenarios poses persistent challenges, such as risk-sensitive exploration [13] and time-consuming episode collection [27] . Recent advances view offline RL as a hopeful solution to these challenges [34] . Offline RL aims to learn a policy from a fixed dataset without further interactions [32] . It can tap into existing large-scale datasets for safe and efficient learning [23, 37, 50] .In offline RL research, a well-known concern is the out-of-distribution (OOD) action issue: the evaluation of OOD actions causes extrapolation error [12] , which can be exacerbated by bootstrapping and result in severe value overestimation [34] . To address this issue, a large body of work has emerged to directly or indirectly suppress OOD actions during training, employing various techniques such as policy constraint [12, 30, 10] , value penalization [31, 2, 6] , and in-sample learning [29, 14, 71] .Distinguished from most previous works, this paper argues that, apart from the OOD action issue, there exists an OOD state issue that also impairs performance yet has received limited attention in the field. Such an issue refers to the agent encountering states out of the offline dataset during the policy deployment phase (i.e., test phase). The occurrence of OOD states can be attributed to OOD actions, stochastic environments, and real-world perturbations. Since typical offline RL algorithms do not involve policy training in OOD states, the agent tends to behave in an uncontrolled manner once entering OOD states in the test phase. This can further exacerbate the state deviation from the offline dataset and lead to severe degradation in performance [34, 75] .In mitigating this OOD state issue, existing limited work attempts to train the policy to correct the agent from OOD states to in-distribution (ID) states [75, 22] . Technically, Zhang et al. [75] Jiang et al. [22] resort to an inverse dynamics model for policy constraint. However, they deal with the OOD state and OOD action issues separately, requiring extra OOD action suppression components and complex distribution modeling, which sacrifices computational efficiency and algorithmic simplicity. Moreover, correcting the agent to all ID states impartially could be problematic, especially when the dataset contains substantial suboptimal states. As a result, the performance of prior methods also leaves considerable room for improvement.In this paper, we aim to address these two fundamental OOD issues simultaneously by proposing a simple yet effective approach for offline RL. We term our method SCAS due to its integration of OOD State Correction and OOD Action Suppression. We start with solving an analytical form of a value-aware state transition distribution, which is within the dataset support but skewed toward highvalue states. Then, we align it with the dynamics induced by the trained policy on perturbed states via KL divergence. This operation intends to correct the agent from OOD states to high-value ID states, a concept we refer to as value-aware OOD state correction. Through some derivations, it also eliminates the necessity of training a multi-modal state transition model. Furthermore, we show theoretically and empirically that, while designed for OOD state correction, SCAS regularization also exhibits the effect of OOD action suppression. We evaluate SCAS on the offline RL benchmarks including D4RL [9] and NeoRL [49] . SCAS achieves excellent performance with consistent hyperparameters without additional tuning. Moreover, benefiting from its OOD state correction ability, SCAS demonstrates improved robustness against environmental perturbations.To summarize, the main contributions of this work are:• We systematically analyze the underexplored OOD state issue in offline RL and propose a simple yet effective approach SCAS unifying OOD state correction and OOD action suppression. • Our approach achieves value-aware OOD state correction, which circumvents modeling complex distributions and significantly improves performance over vanilla OOD state correction methods. • Empirically1 , our approach demonstrates superior performance on standard offline RL benchmarks and enhanced robustness in perturbed environments without additional hyperparameter tuning. 
preliminaries: 
 In reinforcement learning, we generally characterize the environment as a Markov Decision Process (MDP) M = (S, A, P, R, γ, d 0 ), with state space S, action space A, transition dynamics P : S × A → ∆(S), reward function R : S × A → R, discount factor γ ∈ [0, 1), and initial state distribution d 0 [61] . The agent interacts with the environment and seeks a policy π : S → ∆(A) to maximize the expected discounted return η(π):EQUATIONFor any policy π, we define the value function as V π (s) = E π [ ∞ t=0 γ t R(s t , a t )|s 0 = s] and the state-action value function (Q-value function) as Q π (s, a) = E π [ ∞ t=0 γ t R(s t , a t )|s 0 = s, a 0 = a].Offline RL. In offline RL, the agent can only access a static dataset D = {(s i t , a i t , s i t+1 , r i t )}. We denote the empirical behavior policy of D by β(a|s) and the empirical dynamics model by M (s ′ |s, a), both of which depict the conditional distributions observed in the dataset [12] . Typical actor-critic algorithms [56, 18] evaluate policy π by minimizing Bellman loss:EQUATIONwhere π ϕ and Q θ are the parameterized policy and Q function, and Q θ ′ is a target network whose parameters are updated via Polyak averaging [42] .Simultaneously, policy improvement in policy iteration is achieved via maximizing the Q-value:EQUATIONOOD action issue. In offline RL, OOD actions refer to actions outside the support of the behavior policy β(•|s) at a specific state s ∈ D. Since the Q-values of OOD actions can be poorly estimated and the policy improvement is towards maximizing the estimated Q θ , the resulting policy tends to prioritize the OOD actions with overestimated values, leading to poor performance [12] . 
ood state correction: 
 The following focuses on the OOD state issue and OOD state correction in offline RL. In Section 3.1, we systematically analyze the OOD state issue, introduce the concept of OOD state correction, and point out limitations of prior methods. Then we present the proposed approach SCAS in Section 3.2.In offline RL, OOD states refer to states not in the offline dataset. The OOD state issue (Definition 1) pertains to scenarios where the agent enters OOD states during the test phase, potentially resulting in catastrophic failure [34] . However, such a topic is rarely investigated in the literature, and existing studies lack deep insights. We mathematically formulate the OOD state issue as follows.Definition 1 (OOD state issue). There exists s ∈ S, such that d π M T (s) > 0 and d D (s) = 0, where M T is the MDP of the test environment, π is any learned policy, d π M T is the state probability density induced by π in M T , and d D is the state probability density in the offline dataset.Origins and consequence of OOD states. During the test phase, the OOD states occur primarily in three scenarios: (i) OOD actions: the learned policy, not perfectly constrained within the support of the behavior policy, executes unreliable OOD actions, leading to OOD states. (ii) Stochastic environment: the initial state of the actual environment may fall outside the offline dataset. In addition, stochastic dynamics can also lead to states outside the dataset, even when taking ID actions in ID states. (iii) Perturbations: commonly seen in real-world robot applications, some unexpected perturbations can propel the agent into OOD states (e.g., wind, human interference).During offline training, the typical Bellman updates involve only ID states, and the policies in OOD states are not trained. As a result, when encountering OOD states in the test phase, the agent would exhibit uncontrolled behavior, and the state deviation from the offline dataset can be further exacerbated over time steps, severely degrading performance [34] .OOD state correction. To mitigate this OOD state issue, an intuitive solution is to train a policy capable of correcting the agent from OOD states to ID states, a concept known as OOD state correction [75] . Specifically, during offline training, we can perturb the original state s in the dataset into ŝ to generate substantial OOD states. Then consider the scenario where the agent starts from ŝ, follows the trained policy π, and transitions to the next state ŝ′ . To reduce state deviation, ŝ′ is expected to be close to the offline dataset. Thus we can align the distribution of ŝ′ with an ID state distribution to regularize the policy and achieve OOD state correction.Continuing the above train of thought, SDC [75] generates the ID state distribution by feeding the original state s into a trained state transition model N (s ′ |s) of the dataset. This model characterizes the conditional state transition distribution in the dataset and is implemented by a conditional variational auto-encoder (CVAE) [58] . After pretraining a dynamics model M (s ′ |s, a) and the state transition model N (s ′ |s), SDC introduces the following policy regularizer for OOD state correction:EQUATIONwhere ŝ is a Gaussian noise perturbed version of the original state s, σ is the standard deviation of the Gaussian,M (•|ŝ, π(•|ŝ)) is shorthand for E â∼π(•|ŝ) M (•|ŝ, â), and MMD is the maximum mean discrepancy measure. More recently, OSR [22] directly aligns the trained policy distribution at the perturbed state with a CVAE inverse dynamics model to constrain the policy in OOD states.Limitations. However, the regularizers of prior methods are only designed to deal with this OOD state issue. To mitigate OOD actions, they require an additional conservative Q learning (CQL) term [31] in value estimation to penalize Q-values of OOD actions. In addition, the state transition distribution and the inverse dynamics distribution are multi-modal in many scenarios [43] . The necessity of extra OOD action suppression components and complex distribution modeling compromises their computational efficiency and algorithmic simplicity. Moreover, correcting the agent to all ID states impartially could be problematic, particularly when the offline dataset contains a large portion of suboptimal states. In such cases, vanilla OOD state correction can lead to suboptimal behaviors. Consequently, there is also significant potential for improvement in the performance of prior methods.For a more comprehensive discussion of related work, please refer to Appendix A.The objective of this work is to formulate a simple yet effective policy regularizer for offline RL that unifies OOD state correction and OOD action suppression. Moreover, we aim to achieve value-aware OOD state correction, involving the correction of the agent from OOD states to high-value ID states.Value-aware state transition. For the ID state distribution to which the agent is corrected, we expect a value-aware state transition distribution N * (•|s) that lies within the support of the dataset state transition distribution N (•|s) but is skewed toward high-value states s ′ . To ensure stability and, more importantly, to enable our subsequently designed algorithm to circumvent modeling complex distributions, we seek a soft optimal version of it. To this end, we consider the following problem2 :EQUATIONwhere α is a hyperparameter to balance the two terms.The optimization problem above has a closed-form solution:EQUATIONwhere Z(s) = s ′ exp (αV (s ′ )) N (s ′ |s)EQUATIONBy substituting the analytical solution of N * from Eq. ( 6) into the KL divergence, we haveargmin π D KL (N * (•|s)∥M (•|ŝ, π(•|ŝ))) = argmax π E s ′ ∼N (•|s) exp (αV (s ′ )) Z(s) log M (s ′ |ŝ, π(•|ŝ)) .Note that N is the state transition distribution in the dataset, and s ∼ D, s ′ ∼ N (•|s) is equivalent to (s, s ′ ) ∼ D. Thus minimizing Eq. ( 7) is equivalent to maximizing following regularizer:EQUATIONAs a result, R(π) effectively eliminates the need for a pre-trained multi-modal state transition model (N or N * ) and enables direct sampling from the dataset for optimization.However, the normalization factor Z(s) in R(π) can be challenging to compute. We note that the regularizer R(π) is derived from the minimization of the KL divergence in Eq. ( 7). Since we aim to minimize this KL at every state s in D and Z(s) only affects the relative weights at different s, it matters less to precisely restore the correct state weights in D by computing Z(s), which is empirically hard to estimate and may bring more instability. Thus, we replace Z(s) in R(π) with an empirical normalizer exp(αV (s)) for computational stability:EQUATIONWe provide further rationale behind this choice of the empirical normalizer in Appendix C.1.Tractable optimization. Now we shift focus to the optimization of R 1 (π). The expectation with respect to π can be moved outside the logarithm by Jensen's inequality:EQUATIONwhere the equality holds when π is deterministic. In general, it is convenient to maximize the lower bound in Eq. (10) using the reparameterization trick. However, to ensure the equality case in Eq. (10) , we opt to train a deterministic policy π. In this case, we can directly maximize R 1 (π) by computing the gradient of π using automatic differentiation [46] .In contrast to model-based RL methods that typically use the learned dynamics model to roll out multi-step trajectories for policy training [20, 73] , our algorithm utilizes the dynamics model to propagate the gradient of policy and regularize policy training, resulting in significantly enhanced computational efficiency. Moreover, the nature of one-step dynamics prediction in our method is advantageous for maintaining relatively high prediction accuracy. 
analysis of ood action suppression: 
 This section focuses on the OOD action issue and shows that the proposed regularizer also exhibits the effect of OOD action suppression. In other words, it can also prevent the policy from taking OOD actions, thereby simultaneously addressing the fundamental OOD action issue in offline RL. In offline RL, OOD actions are exclusively defined on ID states. This is because actor-critic training is limited to ID states, and any actions on OOD states would not affect training and cause the OOD action issue mentioned in Section 2. Consequently, for the analysis of OOD actions, it is essential to consider ID states. We define R, R1 as the ID state version of R, R 1 , where ŝ = s. R and R1 can be regarded as special cases of R and R 1 , when ŝ sampled from N (s, σ 2 ) is equal to s:EQUATIONEQUATIONThe proposed regularizer functions as follows: when the agent encounters OOD states, it drives the agent to choose actions leading to ID states, as discussed in Section 3.2. When the agent is in ID states, the ID state part of it comes into play. In the following, we show that it helps circumvent taking OOD actions by analyzing the maximizer of R, R1 in tabular MDPs. Proposition 1. Suppose that the environment dynamics is deterministic, then both R(π) and R1 (π) achieve their global maximum at the policy π * , where3 EQUATIONThe support of π * is within that of the behavior policy β:EQUATION)and π * makes the following equation hold:EQUATIONUnder the deterministic dynamics condition, Proposition 1 shows that π * is constrained within the support of the behavior policy. Thus, our regularizer helps to keep the policy from taking OOD actions. Moreover, π * is able to exactly align M (•|s, π * (•|s)) with N * (•|s), indicating the guidance of the agent to the high-value ID state distributions.Furthermore, we show in Proposition 2 that even under stochastic dynamics, the optimization of R and R1 still yields policies constrained within the support of β. Hence, SCAS also exhibits the effect of OOD action suppression in this more general scenario. Proposition 2. When the dynamics is stochastic, the maximizers of both R(π) and R1 (π) are constrained within the support of the behavior policy:EQUATION5 Implementation Details Update ϕ by maximizing J π (ϕ) in Eq. ( 20)10: Update target network: θ ′ ← (1 -τ )θ ′ + τ θ 11:end for SCAS is easy to implement and we design the practical algorithm to be as simple as possible, retaining algorithmic simplicity and improving computational efficiency.Dynamics model. We employ a deterministic dynamics model M ω . The loss for training the model isEQUATIONPolicy improvement. With a deterministic model, we replace the log-likelihood in R 1 (π) with mean squared error. It is a common approach in RL algorithms to convert a maximum likelihood estimation problem into a regression problem when dealing with Gaussians with fixed variance [10] . As discussed in Section 3.2, we also adopt a deterministic policy model π ϕ . Thus, we have the following policy regularizer:EQUATIONwhere V θ (s) = Q θ (s, πϕ (s)) and πϕ means π ϕ with detached gradients. Using deterministic policy also simplifies the training process without learning a V -function. Combining R 2 (π ϕ ) with the standard policy improvement objective, we update the policy by maximizing:EQUATIONwhere λ is a hyperparameter to balance the two terms. Additionally, following TD3+BC [10] , we also normalize Q θ in the first term in each mini-batch to maintain a balanced scale across tasks.Overall algorithm. Putting everything together, we present our final algorithm in Algorithm 1. 
experiments: 
 In this section, we conduct several experiments to examine the performance and properties of SCAS. Please refer to Appendices D and E for experimental details and additional results.OOD state correction. To examine the OOD state correction ability, we compare the state distributions generated by the learned policies of different algorithms with the state distribution of the offline dataset. In detail, we first train SCAS, CQL [31] , and TD3+BC [10] , and then collect 50,000 samples by running the trained policies separately. We also sample 50,000 states randomly from the offline dataset for comparison. Figures 1(a ) to 1(c) plot the state distributions in halfcheetah-mediumexpert [9] with t-SNE [62] , and Figure 1 (d) visualizes the optimal value of each state. We access these values from the learned value function obtained by running TD3 [11] online to convergence.In Figures 1(a ) and 1(b), we observe that the policies learned by CQL and TD3+BC tend to produce OOD states. As depicted in Figure 1 (d), these OOD states have extremely low values, so entering them can be detrimental to performance. In contrast, the state distribution induced by SCAS is almost entirely within the support of the offline distribution, demonstrating the OOD state correction ability of SCAS. Moreover, we also note that in the low-value area of the offline state distribution (the grey circle in Figure 1 (d)), SCAS exhibits a very low state density, which could be attributed to SCAS's value-aware OOD state correction. We refer the reader to Appendix E.2 for additional experiments validating the OOD state correction effects. OOD action suppression. We empirically evaluate the OOD action suppression effects through the lens of value estimates. We compare SCAS with three baselines: (1) ordinary offpolicy RL which is SCAS with λ = 0 (all other implementations are the same); (2) SDC [75] without additional CQL [31] term to suppress OOD actions; (3) OSR [22] without additional CQL term. We conduct experiments on D4RL datasets [9] . Since value over-estimation (divergence) is the main consequence and evidence of OOD actions [12] , we plot the learned Q-values of SCAS and the baselines in Figure 2 .We also include the oracle Q-values of SCAS by rollouting the trained policy for 1, 000 episodes and evaluating the Monte-Carlo return. Additional results are provided in Appendix E.1.The results show that the learned Q-values of ordinary off-policy RL, SDC without CQL, and OSR without CQL diverge at early learning stages, suggesting that the algorithms suffer from severe OOD actions. By contrast, the learned Q-values of SCAS stay close to the oracle Q-values. This indicates that SCAS regularization alone is able to suppress OOD actions.Tasks. We evaluate SCAS on D4RL [9] and NeoRL [49] benchmarks. In D4RL, we conduct experiments on Gym locomotion tasks and much more challenging AntMaze tasks. Due to the space limit, the results on NeoRL are deferred to Table 4 in Appendix E.3. Baselines. We compare SCAS with prior state-of-the-art offline RL methods as well as the ones specifically designed for OOD state correction, including BC [48] , MOPO [73] , OneStep RL [5] , CQL [31] , TD3+BC [10] , IQL [29] , SDC [75] and OSR [22] .Hyperparamter tuning. Offline RL methods are appealing for their ability to generate effective policies without online interaction. Nevertheless, many existing offline RL works involve datasetspecific hyperparameter tuning. The reduction of hyperparameter tuning is crucial for improving practical applicability. In this work, SCAS uses a single set of hyperparameters for all datasets in D4RL and NeoRL benchmarks to obtain the reported results.Comparisons with baselines. On D4RL, comparisons of performance, runtime, and hyperparameter tuning information are shown in Table 1 . We refer the reader to Appendix E.8 for learning curve details of SCAS. On the Gym locomotion tasks, SCAS outperforms prior methods on most datasets and achieves the highest total score with a single set of hyperparameters. On the challenging AntMaze tasks, SCAS performs better than IQL and outperforms other methods by a very large margin. In NeoRL (Table 4 ), SCAS performs comparably to MOBILE [59] and outperforms other baselines.Runtime. We present the runtime of algorithms at the bottom of Table 1 . SCAS exhibits significantly lower runtime than MOPO, SDC, and OSR and is comparable to other model-free baselines.Generality. SCAS is a generic model-based regularizer that can be easily integrated into existing offline RL algorithms. The corresponding results and analysis are provided in Appendix E.5.In this section, we evaluate the algorithms in a more real-world setting where the agent receives uncertain perturbations during test time. OOD state correction is even more critical in such scenarios since the agent can enter OOD states after perturbation. To simulate this scenario, we add varying steps of Gaussian noise with a magnitude of 0.5 to the actions conducted by the policy during test time. Specifically, the policy is trained on standard D4RL datasets but is tested in the perturbed environments. We control the strength of perturbations by adjusting the number of perturbation steps. Figure 3 shows the results of TD3+BC, CQL, SDC, and SCAS on various datasets over five random seeds. We observe that SCAS consistently outperforms previous methods across different perturbation levels and also exhibits less performance degradation against perturbations. Therefore, SCAS enjoys better robustness against perturbations in the complex and unpredictable environments.We examine the effects of the inverse temperature α, the balance coefficient λ, and the noise scale σ.Due to the space limit, the results for σ and on additional datasets are deferred to Appendix E.6. A sensitivity analysis on dynamics model errors is also provided in Appendix E.7.Inverse temperature α. α is the key hyperparameter in SCAS for achieving value-aware OOD state correction. If α = 0, the effect degenerates to vanilla OOD state correction. Figure 4 (a) displays the learning curves of SCAS with different α. The results show that a large α is crucial for achieving good performance (also verified on more tasks in Figure 6 ), clearly demonstrating the effectiveness of our value-aware OOD state correction. However, too large α (α = 10) induces less satisfying performance, probably due to the increased variance of the learning objective.Balance coefficient λ. λ in Eq. ( 20) controls the balance between vanilla policy improvement and SCAS regularization. We vary λ within the range [0, 1] and present the learning curves of SCAS in Figure 4 (b). Notably, SCAS is able to converge to good performance over a very wide range of λ (also verified on more tasks in Figure 7 ). An interesting finding is that even when λ = 1 and the signal from RL improvement (max Q) is removed, SCAS still performs well on most tasks. This could be attributed to the fact that value-aware OOD state correction implies some sort of improvement in policy by maximizing the values of policy-induced next states.Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?Answer: [Yes] Justification: Please refer to Appendix D.Guidelines:• The answer NA means that the paper does not include experiments.• The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them. • The full details can be provided either with the code, in appendix, or as supplemental material. 
conclusion and limitations: 
 In this paper, we systematically analyze the OOD state issue in offline RL and propose SCAS, a simple yet effective approach that unifies OOD state correction and OOD action suppression. SCAS also achieves value-aware OOD state correction, significantly improving performance over vanilla OOD state correction. Empirical results validate the properties of SCAS, showcasing its superior performance on the offline RL benchmarks and its enhanced robustness in perturbed environments.However, our work also has some limitations. For example, current SCAS primarily focuses on continuous control tasks. In discrete settings, algorithmic components like state perturbation strategy would be different, which would be an interesting direction for future work. Moreover, we anticipate employing more advanced dynamics models, such as ensembles [73] and diffusion models [21] , to further improve the performance of our method.Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?Answer: [Yes]Justification: The results in the paper are accompanied by standard deviations across multiple seeds.Guidelines:• The answer NA means that the paper does not include experiments.• The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper. • The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions). • The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.) • The assumptions made should be given (e.g., Normally distributed errors).• It should be clear whether the error bar is the standard deviation or the standard error of the mean. • It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified. • For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates). • If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text. 
experiments compute resources: 
 Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?Answer: [Yes] Justification: Please refer to Appendix D.Guidelines:• The answer NA means that the paper does not include experiments.• The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage. • The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute. • The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper). 
code of ethics: 
 Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines? 
broader impacts: 
 Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?Answer: [Yes] Justification: Please refer to Appendix F.Guidelines:• The answer NA means that there is no societal impact of the work performed.• If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact. • Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.• The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster. • The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology. • If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?Answer: [NA] Justification: The paper poses no such risks.Guidelines:• The answer NA means that the paper poses no such risks.• Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters. • Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images. • We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.12. Licenses for existing assets Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?Answer: [Yes] Justification: The creators or original owners of assets used in the paper are properly credited and the license and terms of use are explicitly mentioned and properly respected.Guidelines:• The answer NA means that the paper does not use existing assets.• The authors should cite the original paper that produced the code package or dataset.• The authors should state which version of the asset is used and, if possible, include a URL. • The name of the license (e.g., CC-BY 4.0) should be included for each asset.• For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided. • If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset. • For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided. 
 
 Appendix: Model-free offline RL. In offline RL, extrapolation error and overestimation caused by OOD actions pose significant challenges. Among model-free solutions, value regularization methods penalize the Q-values of OOD actions [31, 2, 28, 36, 3, 72, 40] , while policy constraint approaches compel the trained policy to be close to the behavior policy, either explicitly via divergence penalties [69, 30, 10] , implicitly by weighted behavior cloning [47, 45, 67, 39] , or directly through specific parameterization of the policy [12, 15] . Relatively independently, in-sample learning methods formulate the Bellman target using only the actions in the dataset to avoid OOD actions [5, 29, 76, 71] . Recently, some works aim to learn the optimal policy within the support of the dataset (known as in-support or in-sample optimal policy) in a theoretically sound way and are less affected by the average quality of the dataset [39, 40, 68] . However, existing popular offline RL approaches primarily focus on the OOD action issue during training and often neglect the OOD state issue during the test phase.Model-based offline RL. Model-based RL methods learn a model of the environment and generate synthetic data from that model to optimize the policy [60, 20, 24] . To ensure conservatism in offline RL, Kidambi et al. [25] and Yu et al. [73] estimate the uncertainty in the model and apply reward penalties for state-action pairs with high uncertainty. Some model-based approaches also introduce conservatism similarly to model-free ones, employing techniques like value regularization [74] and policy constraint [41] . Recently, Sun et al. [59] conducts uncertainty quantification through the inconsistency of Bellman estimations under the learned dynamics ensemble. However, model-based methods often come with a high computational burden [20] , and their effectiveness relies heavily on the quality of the trained model [43] . In contrast, our algorithm leverages the dynamics model to propagate policy gradients, make one-step predictions, and regularize policy training, leading to significantly improved computational efficiency and relatively high prediction accuracy.OOD state correction. In offline RL, OOD state correction deserves more attention as the state deviation during the test phase can accumulate over time steps, severely degrading performance [34] .Existing limited solutions aim to train the policy to correct the agent from OOD states to ID states [75, 22] . Specifically, SDC [75] builds a dynamics model and a state transition model, and aligns the policy-induced next state distributions at OOD states with the state transition model. On the other hand, OSR [22] utilizes an inverse dynamics model to constrain the policy at OOD states. Compared with prior methods, our proposed SCAS efficiently unifies OOD state correction and OOD action suppression in offline RL and additionally achieves value-aware OOD state correction. The DICE series of works [44, 33, 38] share similar motivations with SCAS to some extent; however, there are significant differences between the two. Firstly, DICE is based on a linear programming framework of RL, while SCAS is based on a dynamic programming framework. Therefore, the theoretical foundations and learning paradigms of the two are inherently different. Secondly, SCAS only corrects encountered OOD states, whereas DICE algorithms require the policy-induced occupancy distribution to align with the dataset distribution. Therefore, DICE's constraints are stricter, potentially making it more susceptible to the average quality of the dataset. Lastly, theoretical and empirical evidence indicate that DICE algorithms have a problem of gradient cancellation [38] , which imposes certain limitations on their practical effectiveness.In this section, we present the proofs for the theories in the paper.We show that Eq. ( 6) is the optimal solution of the optimization problem Eq. ( 5):EQUATIONWe can optimize N * at each s ∈ D separately. Thus we consider the following optimization problem:EQUATIONThis constrained optimization problem is convex, and the Lagrangian is:EQUATIONThe KKT condition gives:EQUATIONSolving for Ñ gives the closed form solution N * :EQUATIONBy the condition s ′ N * (s ′ |s) = 1, we can directly solve the Lagrangian multiplier ν and replace exp(ν -1) with a normalization factor:EQUATIONwhere Z(s) = s ′ exp (αV (s ′ )) N (s ′ |s) is the normalization factor.The support of π * is within that of the behavior policy β:EQUATIONand π * makes the following equation hold:EQUATIONProof. We start with R(π).EQUATIONEQUATION= argminEQUATION= argminEQUATIONThe third equality holds because of the relationship between N * and N in Eq. ( 6):EQUATIONTherefore, the maximizer of R(π) is equal to the solution of the minimization problem in Eq. (35) . Now consider the two distributions N * (•|s) and M (•|s, π(•|s)) in Eq. (35) .EQUATIONFor analytical clarity, we use the notation M with slightly different meanings in different cases: in the stochastic setting, M : S × A → ∆(S); in the deterministic setting, M : S × A → S. With the deterministic dynamics assumption,EQUATIONOn the other hand, EQUATIONEQUATIONWe first show that π * is a valid policy, that is, π * is normalized.EQUATION= exp (αV (M (s, a))) β(a|s)EQUATION= exp (αV (M (s, a))) β(a|s)EQUATION= exp (αV (M (s, a))) β(a|s) EQUATIONTherefore, a π * (a|s) = 1.Substitute Eq. ( 44) into Eq. ( 41),EQUATIONComparing Eq. ( 43) with Eq. ( 50), it holds that N * (s ′ |s) = M (s ′ |s, π * (•|s)), ∀s ∼ D. As a result,EQUATIONConsidering the non-negativity of KL divergence, the optimization problem in Eq. ( 35) achieves its global minimum at π * . Therefore, R(π) also achieves its global maximum at π * .EQUATIONEQUATION= argmaxEQUATION= argminEQUATIONEQUATIONThe fourth equality holds because of the relationship between N * and N in Eq. ( 36).As shown above, it holds that N * (sEQUATIONConsidering Z(s)/exp (αV (s)) > 0 and the non-negativity of KL divergence, the optimization problem in Eq. ( 58) achieves its global minimum at π * . Therefore, R1 (π) also achieves its global maximum at π * .In conclusion, when the environment dynamics is deterministic, both R(π) and R1 (π) achieve their global maximum at the policy π * , and π * makes the following equation hold: Proposition 4 (Proposition 2 in the main paper). When the dynamics is stochastic, the maximizers of both R(π) and R1 (π) are constrained within the support of the behavior policy:EQUATIONEQUATIONProof. We start with R(π).EQUATIONLet π denote any valid policy. For ∀s ∈ D, define ϵ(s) and n(s) as follows:EQUATIONEQUATIONFor ∀s ∈ D, there exists at least one action a such that (s, a) ∈ D. Thus it holds that n(s) > 0, ∀s ∈ D. Then for ∀s ∈ D, ∀π, define π in as follows:π in (a|s) = π(a|s) + ϵ(s) n(s) , β(a|s) > 0, 0, β(a|s) = 0.π in can be seen as a projection of π onto β's support. Besides, for ∀s ∈ D, = 1Thus π in is a valid policy.EQUATIONEQUATIONEQUATION≥ 0 (78)The second equality holds because, in tabular MDPs, the empirical dynamics model M exactly computes the conditional distribution observed in the dataset. For transitions not contained in the dataset, M = 0 [12] . The final inequality holds because ϵ(s) ≥ 0.EQUATION≥ 0 (82) Now suppose π is not constrained within the support of the behavior policy at some state s 1 ∈ D: supp(π(•|s 1 )) ⊈ supp(β(•|s 1 )). That is, ∃ã 1 such that π(ã 1 |s 1 ) > 0 and β(ã 1 |s 1 ) = 0. Thus it holds that ϵ(s 1 ) = a I[β(a|s 1 ) = 0]π(a|s 1 ) > 0. On the other hand, since s 1 ∈ D, there exists at least one action a 1 and one state s ′ 1 such that (s 1 , a 1 , s ′ 1 ) ∈ D. Thus it holds that β(a 1 |s 1 ) > 0 andEQUATION> 0In such case, R(π in ) > R(π). Therefore, if π is not constrained within the support of the behavior policy at some state s 1 ∈ D, we can find another policy π in that is constrained within the support of the behavior policy and achieves higher objective function R(π in ). Consequently, R(π) must achieve its maximum at support constrained policy π * : supp(π * (•|s)) ⊆ supp(β(•|s)), ∀s ∼ D.EQUATIONWith the same definition of ϵ(s), n(s) and π in as in Eq. ( 66), Eq. ( 67) and Eq. ( 68), it also holds thatR1 (π in ) -R1 (π) (88) = E (s,s ′ )∼D exp (α (V (s ′ ) -V (s))) log a M (s ′ |s, a)π in (a|s) a M (s ′ |s, a)π(a|s) (89) ≥ E (s,s ′ )∼D [exp (α (V (s ′ ) -V (s))) log (1)] (90) ≥ 0 (91)As before, when supposing π is not constrained within the support of the behavior policy at some state s 1 ∈ D, it holds that R1 (π in ) > R1 (π). Therefore, R1 (π) must achieve its maximum at support constrained policy π * 1 : supp(π * 1 (•|s)) ⊆ supp(β(•|s)), ∀s ∼ D. In conclusion, when the environment dynamics is stochastic, the maximizers of both R(π) and R1 (π) are constrained within the support of the behavior policy:EQUATIONC Further Discussions C.1 Rationale for Choosing exp(αV (s)) as the Empirical NormalizerFirstly, choosing exp(αV (s)) is intended to obtain something similar to the advantage function.With this normalizer, the weight of our regularizer is exp(α(V (s ′ ) -V (s))), which is comparable to the weight exp(αA(s, a)) in Advantage Weighted Regression (AWR) [47] . Here, V (s ′ ) -V (s) represents the relative advantage of the next state s ′ compared to the current state s, while A(s, a) reflects the relative advantage of taking action a in s compared to following the current policy.Comparison of the objectives of SCAS and AWR:EQUATIONAWR: exp(αA(s, a)) log π(a|s)Secondly, as discussed in the paper, introducing any normalizer that depends only on s (independent of s ′ ) does not affect the development and analysis of our method; it is merely for computational stability. In AWR-based methods, there also exists a normalizer Z(s) and they usually disregard it [47, 45] . The rationale behind this is similar.In a specific sense, SCAS, which unifies OOD state correction and OOD action suppression, also integrates pessimism and state robustness. (1) Regarding pessimism: The OOD action suppression effect of SCAS aligns with the pessimism commonly discussed in offline RL work (being pessimistic about OOD actions) [31, 70, 30, 3, 54] . Unlike traditional policy constraint methods [69, 30, 10, 47] , our approach does not require the training policy to align with the behavior policy; it only requires the successor states to be within the dataset support, which is a more relaxed constraint. ( 2) Regarding state robustness: The OOD state correction effect of SCAS is aimed at improving the agent's robustness to OOD states during the test phase. Compared with previous works, SCAS unifies OOD state correction and OOD action suppression and additionally achieves value-aware OOD state correction. Some offline RL literature on state robustness differs from our approach; they typically consider noisy observations [72] , such as sensor errors. In contrast, SCAS addresses state robustness concerning actual OOD states encountered during test time, rather than noisy observations.In SCAS, there is regularization on the policy's output actions at ID states. In our regularizer, the perturbed states ŝ are sampled from N (s, σ 2 ), and a large portion of ŝ will fall near the original ID state s or even be approximately equal to s. Therefore, the policy's output actions at ID states are also regularized. For this part of the regularization, its role is equivalent to the ID state regularizer analyzed in Section 4, which has been theoretically shown to have the effect of suppressing OOD actions. Moreover, the experimental results in Section 6 also demonstrate that our OOD state correction regularizer addresses the traditional issue with OOD actions.We further elucidate the differences between the well-known OOD action issue and the OOD state issue we analyzed. Most offline RL works focus on the OOD action issue in the training phase. That is, the trained policy outputs OOD actions to compute the target Q, which results in extrapolation error and value divergence during training [12] . In contrast, the OOD state issue we defined and analyzed is in the test phase. That is, the agent can enter states out of the offline dataset during test, potentially resulting in catastrophic failure. 2 . Note that we use this same set of hyperparameters to obtain all the results reported in this paper (except for parameter study). Following TD3+BC [10] , we normalize the states in all datasets except for antmaze-large. We clip the exponentiated weight exp (αV θ (s ′ ) -αV θ (s)) in Eq. ( 19) to (-∞, 50]. Following the suggestions in the benchmark [9] , we subtract 1 from the rewards for the Antmaze datasets.Our evaluation criteria follow those used in most previous works. For the Gym locomotion tasks, we average returns over 10 evaluation trajectories and 5 random seeds, while for the Ant Maze tasks, we average over 100 evaluation trajectories and 5 random seeds. The reported results are the normalized scores, which are offered by the D4RL benchmark [9] to measure how the learned policy compared with random and expert policy:D4RL score = 100 × learned policy returnrandom policy return expert policy returnrandom policy returnThe results of baselines reported in Table 1 are obtained as follows. We re-run OSR [22] on all datasets using their official codebase 5 and tune the hyperparameters for each dataset as specified in their paper. We implement SDC [75] and re-run it on all datasets. We use the SDC-related hyperparameters as specified in their paper, and sweep the CQL-related hyperparameters in {1,2,5,10,20} for each dataset. We re-run OneStep RL [5] on all datasets using their official codebase 6 and the default hyperparameters. We implement BC [48] based on the TD3+BC repository 7 and re-run it on all datasets. The results of other baselines are taken from [3] and [68] . The runtime in Table 1 is obtained by running offline RL algorithms on halfcheetah-medium-replay-v2 on a GeForce RTX 3090.Figures 1(a ) to 1(d) share the same embedding function obtained by running t-SNE on the set of all 200,000 samples (50,000 samples each from the dataset, CQL, TD3+BC, and SCAS). This ensures a clear visual comparison. Figure 1(d ) contains all the 200,000 samples, which is the union of the points in Figures 1(a ) to 1(c).Under the same setting of Figure 2 , we conduct experiments on the additional datasets. The results are shown in Figure 5 . We omit the Q values of Off-policy RL, SDC w/o CQL, and OSR w/o CQL at higher numbers of optimization steps, because these Q values diverge in the early learning stage, and plotting their Q values at later optimization steps would result in an excessive range on the vertical axis. The additional results also show that only SCAS's OOD state correction term can achieve OOD action suppression and prevent value over-estimation. walker2d-medium-expert-v2 Here Off-policy RL is SCAS with weight λ = 0 in Eq. (20) . Only SCAS's OOD state correction term can achieve OOD action suppression and prevent value over-estimation (divergence).To further examine the OOD state correction effects of SCAS, we conduct experiments on a modified D4RL maze2d-open-v0 [9] . It is a 2D point robot navigation task in a rectangle map with vertices (0, 0) and (3, 5) . The agent needs to reach the goal at (2, 3). We modify the dataset by removing all the transitions containing states in a rectangle with vertices (0, 0) and (1.5, 2.5). During test, we let the initial state be randomly distributed in this OOD region. We train algorithms over 10 6 gradient steps and average returns over 1000 evaluation trajectories.The results of BC [48] , TD3+BC [10] , CQL [31] , MOPO [73] , IQL [29] , and SCAS are reported in Table 3 . With the OOD state correction signals, SCAS corrects the agent out of the OOD region more quickly and stably, achieving significantly better performance than typical offline RL methods. We also evaluate SCAS on the NeoRL benchmark [49] . NeoRL is a benchmark designed to simulate real-world scenarios by collecting datasets using a more conservative policy, aligning closely with realistic data collection scenarios. The narrow and limited data makes it challenging for offline RL algorithms. The results are shown in Table 4 . The results of baselines are taken directly from the MOBILE paper [59] . According to Appendix C in [59] , these results are obtained by tuning hyperparameters per dataset. For SCAS, we use the same fixed set of hyperparameters as specified in Appendix D. Without additional hyperparameter tuning, SCAS still performs comparably to MOBILE and outperforms other baselines in total scores.The original SCAS requires only one single hyperparameter configuration in implementations. For a fair comparison with DW [19] , EDAC [2] , RORL [72] , SQL [71] , and EQL [71] , we roughly select λ from {0.025, 0.25} for each dataset, referring to this variant as SCAS-ht. The results of SCAS-ht and these methods are reported in Table 5 . Among the ensemble-free methods, SCAS-ht achieves the highest performance in both mujoco locomotion and antmaze domains. Compared with ensemble-based methods, SCAS-ht also performs better on antmaze tasks. DW [19] reweights ID data points by their values for behavior regularization and does not account for OOD states during the test phase. In contrast, our approach considers an OOD state correction scenario, resulting in enhanced robustness during the test phase and better performance.The SCAS regularizer is compatible with various offline RL objectives. We conduct experiments to combine SCAS with CQL [31] , IQL [29] , and TD3BC [10] . Comparisons between these combined algorithms and the original ones are shown in Table 6 . We find that applying the SCAS regularizer leads to improved performance for these popular algorithms, which could be attributed to the OOD state correction effects of SCAS. However, we also find that these combined methods do not achieve better performance than the original SCAS (comparable on most tasks and worse on some tasks). We hypothesize that this is because SCAS already has the effect of OOD action suppression, and when combined with offline RL objectives that also aim for OOD action suppression, it may become overly conservative. As a result, the combined algorithms may perform worse than the original SCAS on some sub-optimal datasets.In this section, we present additional parameter study results conducted on four challenging Antmaze tasks, including antmaze-large-play-v2, antmaze-large-diverse-v2, antmaze-medium-play-v2, and antmaze-medium-diverse-v2.Inverse Temperature α. The inverse temperature α is the key hyperparameter in SCAS for achieving value-aware OOD state correction. It controls the significance of the values of next states in SCAS's OOD state correction. If α = 0, the effect corresponds to vanilla OOD state correction. As α gets larger, SCAS is more inclined to correct the agent to the high-value ID states. Thus we can assess the effectiveness of value-aware OOD state correction compared to vanilla OOD state correction by varying α. Here we test SCAS with different α and the results are shown in Figure 6 . We observe that a large α is crucial for achieving good performance on all the antmaze tasks, clearly demonstrating the effectiveness of our value-aware OOD state correction. However, too large α (α = 10) induces less satisfying performance, probably due to the increased variance of the learning objective. In general, we find that choosing α = 5 leads to the best performance. Balance Coefficient λ. The balance coefficient λ controls the balance between vanilla policy improvement and SCAS regularization. If we set λ = 0, SCAS degenerates into the vanilla off-policy RL algorithm. Here we vary λ in {0, 0.25, 0.5, 0.75, 1} and present the corresponding learning curves of SCAS in Figure 7 . Notably, SCAS is able to converge to good performance over a very wide range of λ. However, if λ = 0, the vanilla off-policy RL suffers from extrapolation error and overestimation, demonstrating poor performance. We also observe a very interesting fact that even when λ = 1 and the signal from RL improvement (max Q) is removed, SCAS still performs well on most tasks. This could be attributed to the fact that value-aware OOD state correction implies some sort of improvement in policy by maximizing the values of policy-induced next states. Noise Scale σ. The noise scale σ is the standard deviation of the Gaussian noise added to the original states for formulating the SCAS regularizer. Here we test SCAS with different σ and present the corresponding learning curves in Figure 8 . We observe a significant performance drop with too large σ (σ = 1) on all the tasks, due to the heavily corrupted learning signal. On the other hand, when σ = 0 (without noise), the performance is also less satisfying. With σ = 0, SCAS is still able to prevent the agent at ID states from entering OOD states, maintaining the agent in safe regions, but it cannot correct the agent from OOD states to ID states as reliably as the original SCAS. In general, we find that choosing σ = 0.001 or 0.01 leads to the best performance. Learning curves on Gym locomotion tasks and Antmaze tasks are presented in Figure 10 and Figure 11 respectively. The curves are averaged over 5 random seeds, with the shaded area representing the standard deviation across seeds.Offline RL holds promise for facilitating practical RL applications in domains like robotics, healthcare, and education, where data collection is often costly or risky. However, it is important to recognize its potential negative societal impacts. One concern is that biases in offline data may transfer to the learned policy. In addition, offline RL may affect employment by automating tasks traditionally Justification: Please refer to the code in the supplemental material. Guidelines:• The answer NA means that paper does not include experiments requiring code. • Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.Justification: The research conducted in the paper conforms, in every respect, with the NeurIPS Code of Ethics.Guidelines:• The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.• If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics. • The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction).Our code is available at https://github.com/maoyixiu/SCAS.Note that the regularizer DKL(N * (•|s)∥N (•|s)) can constrain the support of N * (•|s) within that of N (•|s), because if supp(N * (•|s)) ⊈ supp(N (•|s)) at some state s, then DKL(N * (•|s)∥N (•|s)) = ∞.Here for clarity, we use the notation M with slightly different meanings in different cases: in the stochastic setting, M : S × A → ∆(S); in the deterministic setting, M : S × A → S.Here for clarity, we use the notation M with slightly different meanings in different cases: in the stochastic setting, M : S × A → ∆(S); in the deterministic setting, M : S × A → S.https://github.com/Jack10843/OSRhttps://github.com/davidbrandfonbrener/onestep-rlhttps://github.com/sfujim/TD3_BC 


Decision: All reviewers, including myself as the Area Chair, commend the originality of this submission, particularly its focus on out-of-distribution (OOD) states rather than the more commonly explored OOD actions. The work is grounded in solid theoretical foundations and presents a novel method that is both simple and elegant. While the empirical results only represent an incremental advance, the submission has significant potential to pave the way for future impactful research in the field of Offline Reinforcement Learning. Therefore, I recommend its acceptance.

Reviews: review: paper_summary: paper_summary: The paper proposes a regularization term in offline RL, which can simultaneously address OOD state correction and OOD action suppression without pretraining a state transition model $N(s'|s)$. It shows good performance in D4RL benchmarks.
strength_weakness: strength_weakness: The method utilizes a deterministic policy, which is often regarded as lacking expressiveness. Thus, the performance is not as good as diffusion-based policy methods.
questions: questions: 1. For Definition 1, if $d^\pi_{M_\tau}$ is a continuous probability density function of $s$, and if $d_D(s)$ only has support on $s$ which appears in the dataset, does it mean that any state which doesn't show up in the dataset is OOD? If that's the case, the probability of an OOD state is 1 since the number of ID states must be finite in the dataset?

2. Instead of pretraining $N(s'|s)$, the paper chooses to add Gaussian noise to $s$, with the noise $\sigma=0.003$. Is this hyperparameter important? If $\sigma$ increases, can we have a more robust algorithm since it makes more OOD states into consideration?

3. The method uses a deterministic policy. What if we choose a Gaussian policy?

4. In Figure 2, why do Off-policy, SDC, and OSR only have one dot? Shouldn't it also be a line?

5. How are the results in Table 1 recorded? Final round results, last 5 epoch average, or another method?

6. Where do the baseline method results come from, their original paper? I ask because I found SDC doesn't have results for antmaze tasks. Did the authors implement it by themselves?

7. The authors use antmaze-v2 for their method evaluation, while some baseline methods, for example, IQL and OSR, use antmaze-v0. In my experience, the v2 dataset can always provide higher rewards than v0. Can the authors explain the mismatch? Or can the authors show results on v0 as well?
limitations: limitations: I believe the theory of the paper is solid, while the experimental performance is my concern. I am happy to increase my score if the antmaze version mismatch issue is resolved.
paper_summary: paper_summary: This paper focuses on OOD state issue, an important but overlooked issue in offline RL. This paper proposes aligning the OOD state towards In-Distribution (ID) states with high value, named as value-aware OOD state correction. Additionally, the paper discovers that the overestimation of OOD actions can also be mitigated during the implementation of OOD state correction. The experimental results demonstrate that the proposed algorithm is superior in performance and robustness.
strength_weakness: strength_weakness: 1. The paper proposes that during the process of OOD state correction, the overestimation of OOD actions will also be addressed. I have some doubts about this point. The constraint $R_2$ seems to only regularize the policy to output actions at OOD states that result in the next state being within the offline dataset. However, there is no regularization on the policy's output actions at ID states. Therefore, I am confused about how the paper solves the traditional issue with OOD actions. If I have misunderstood anywhere, please point it out.
questions: questions: 1. Figure 2 suffers from an excessive range on the vertical axis, causing the two curves of SCAS to almost overlap. A suggestion would be to take the log of the Q values on the vertical axis; this should make the two curves more distinguishable.
limitations: limitations: Please see weakness.
paper_summary: paper_summary: This paper presents SCAS (OOD State Correction and Action Supression), a model-based regularization approach that effectively addresses the challenges of out-of-distribution (OOD) states and actions in offline reinforcement learning (RL) algorithms.
The method unfolds in two main stages: (1) training a transition model, and (2) training a policy that incorporates a model-based regularizer. 
This regularizer is designed to steer the policy towards high-value in-distribution (ID) successor states and thereby away from OOD successor states.
While the primary focus of SCAS is to resolve OOD state issues, this paper shows that the regularization strategy not only mitigates visiting OOD states but, as a byproduct, also suppresses OOD actions in the training of the value function.
Empirical results demonstrate that the proposed method performs comparably on standard benchmarks, showing its effectiveness in refining offline RL algorithms.
strength_weakness: strength_weakness: **1. Unclear Motivation for Value-Aware OOD State Correction**

This paper proposes shifting the OOD state distribution not to a standard ID state distribution, but to a high-value ID state distribution instead. This choice raises questions about the specific objectives of value-aware state correction. Is it strategically targeting high-value data points, or is it intended to mitigate the effects of distribution shifts? It seems that the manuscript primarily focuses on the former, in that case, it would benefit from a comparison with existing works like DW[1], which also focus on high-value data points for behavior regularization. If the goal is indeed to mitigate the effects of distribution shifts, then it is crucial for the manuscript to better articulate and emphasize how value-aware OOD state correction contributes to the robustness against OOD state visitation. Moreover, a clearer explanation of how this approach aligns with offline RL principles like pessimism and robustness to OOD states would greatly strengthen the manuscript's validity and impact.

**2. Lack of Comparisons with Comparable Offline RL Approaches**

While the paper tackles a fundamental problem in offline RL, the evaluations presented are somewhat constrained by a lack of comparison with key recent advancements in the field. Notable works such as DW[1], EDAC[2], RORL[3], and SQL/EQL[4], which explore similar offline RL settings, are not considered in the main experiments. This gap is particularly significant as DW[1] employs a value-aware strategy that closely mirrors the approach of this paper, focusing on constraining policies to high-value data points. To convincingly justify the need for correcting OOD states—or for employing a value-aware correction approach—a comparative analysis with these influential studies is essential. This would not only place the findings in a broader context but also potentially highlight the contributions and advantages of the proposed method.

[1] Hong et al., “Beyond Uniform Sampling: Offline Reinforcement Learning with Imbalanced Datasets.”, NeurIPS 2023.

[2] An et al., "Uncertainty-Based Offline Reinforcement Learning with Diversified Q-Ensemble.", NeurIPS 2021.

[3] Yang et al., “RORL: Robust Offline Reinforcement Learning via Conservative Smoothing.”, NeurIPS 2022.

[4] Xu et al., “Offline RL with No OOD Actions: In-Sample Learning via Implicit Value Regularization.”, ICLR 2023.
questions: questions: Q1. How sensitive is the performance to the accuracy of the model?

Q2. Given that a main contribution of this work is a model-based regularizer, it appears that SCAS could be compatible with various offline RL methods. Is it feasible to apply the regularizer across different offline RL objectives, such as CQL, IQL, and others?

Q3. In this study, the trained model is exclusively utilized as a regularizer during policy training. Could this approach be extended to also incorporate the model during testing time?
limitations: limitations: The authors included their limitations in Section 7 and limitations are addressed appropriately.
paper_summary: paper_summary: The issue of OOD actions is a well-known concern in offline RL research. The issue of OOD states, however, is relatively unexplored. In this paper, the authors shed light on the significance of OOD state correction and propose a SCAS algorithm to guide the agent back into high-value ID states when encountering an OOD state. To achieve this goal, SCAS first defines a conditional state transition probability $N^*(s'\mid s)$ skewed towards high-valued states and regularizes the policy by aligning the dynamics induced by the policy on a perturbed state $\hat{s}\sim \mathcal{N}(s, \sigma^2I)$ with the value aware state transition probability $N^*(\cdot \mid s)$. Also, theoretical analysis shows that SCAS implicitly guarantees the correction of OOD actions, thereby removing the need for additional regularizers.
strength_weakness: strength_weakness: 1. The authors did not cite works from the DICE literature(e.g., Nachum et al., 2019; Lee et al., 2021; Mao et al., 2024). Although they do not explicitly correct OOD states, I believe they are deeply related to SCAS because they aim to align the policy's state distribution $d^\pi$ with the dataset's state distribution, thus preventing the policy from visiting OOD states.

2. The choice of $\exp(\alpha V(s))$ as the empirical normalizer seems extremely arbitrary. Providing some rationale behind this choice is recommended.

3. All of the experimental analysis was performed on MuJoCo environments, which have deterministic transition dynamics. Since SCAS utilizes a deterministic environmental model, it may not perform well in environments with stochastic transition dynamics.

4. NeoRL is a relatively uncommon benchmark compared to D4RL. A brief explanation of NeoRL would be helpful for the readers.

5. Some experimental details are missing from the paper.

    * Do Figures 1a to 1d share the same embedding function? How was this embedding function obtained? Was it by running t-SNE on the set of 200,000 samples(50,000 samples each from the dataset, CQL, TD3+BC, and SCAS)?
    * The points in Figure 1d seem to be the union of the points in Figures 1a, 1b, and 1c. Am I correct?
    * Why does Figure 2 omit Q values of off-policy RL, SDC, and OSR for higher numbers of optimization steps?
    * Section 6.3 states that varying steps of Gaussian noise were added to the actions during test time. Figure 3 indicates that the authors added noise up to 40 steps. However, trajectories are usually longer than 40 steps. How were those 40 steps selected? Did the authors add noise to the first 40 interactions with the environment? 

6. Minor comment: The meaning of the term *$\mathcal{D}(s)$ weights* on Line 162 is unclear.


### References

Lee, Jongmin, et al. "Optidice: Offline policy optimization via stationary distribution correction estimation." International Conference on Machine Learning. PMLR, 2021.

Mao, Liyuan, et al. "Odice: Revealing the mystery of distribution correction estimation via orthogonal-gradient update." arXiv preprint arXiv:2402.00348 (2024).

Nachum, Ofir, et al. "Algaedice: Policy gradient from arbitrary experience." arXiv preprint arXiv:1912.02074 (2019).
questions: questions: Please refer to the **Weaknesses** section. As a side note, have you tried using the Q and V functions learned by IQL (or any other value-based offline RL algorithm) instead of learning them together with the policy? From Figure 9, the training process seems a bit unstable, so using a stationary training objective might help.
limitations: limitations: The authors have adequately addressed their work's limitations and potential negative societal impact.

Comments: title: Response to Author Rebuttal
Additional rebuttal contents
Additional rebuttal contents
Additional rebuttal contents
comment: Thank you for your positive feedback! We sincerely appreciate the time and effort you have dedicated to reviewing our manuscript. Your comments and suggestions are highly valuable to us, and we will carefully incorporate them into the latter revision.
Thank you for your positive feedback! We sincerely appreciate the time and effort you have dedicated to reviewing our manuscript. Your comments and suggestions are highly valuable to us, and we will carefully incorporate them into the latter revision.
Thank you for addressing my confusion. Incorporating the description for $R2$ and the refined figures into the future version would further improve the readability of the paper.

Taking into account the supplementary experiments that validate the generalizability of the approach, I have raised my score to support the acceptance of the paper.
Thank you for your insightful comments and for conducting the extensive set of additional experiments!

Since my primary concern regarding the lack of experimental comparisons has been fully addressed, I have raised my rating. 

This work now presents (1) promising empirical results and (2) a generic model-based regularizer that can be easily integrated into existing offline RL algorithms. Hence, I believe that these contributions are substantial for the advancement of the offline RL community.
Thank you for your positive feedback! We sincerely appreciate the time and effort you have dedicated to reviewing our manuscript. Your comments and questions are highly valuable to us.
The reviewer appreciates the authors' detailed response and thorough experiments. I have no further questions at this stage and have updated my score accordingly. I am inclined toward a positive outcome for the paper.
Thank you for your positive feedback! We truly appreciate the time and effort you devoted to reviewing our manuscript. Your comments and suggestions are very helpful.
Thank you for your response. All of my concerns were resolved, so I have updated my score.
------Thank you for continuing to read!------

**Q7: Minor comment: The meaning of the term $\mathcal D(s)$ weights on Line 162 is unclear.**

Thanks for pointing that out. Here, $D(s)$ indicates the offline dataset's empirical state distribution.

**Q8: Have you tried using the Q and V functions learned by IQL (or any other value-based offline RL algorithm) instead of learning them together with the policy?**

Yes, when implementing the algorithm, we experimented with using IQL to train the value function. The results of this variant, IQL+SCAS, are shown in **Table 2** of the PDF. However, we found that IQL+SCAS does not perform better than the original SCAS, which uses vanilla policy evaluation. The performance of IQL+SCAS is comparable on most tasks and slightly worse on some tasks. We have two hypotheses: (1) The policy and value training objectives in IQL+SCAS are not well aligned, which is a known issue with IQL identified in [5]; and (2) SCAS already has the effect of suppressing OOD actions and using IQL to train the value function might introduce additional conservatism.

**Reference**

[1] Lee, Jongmin, et al. "Optidice: Offline policy optimization via stationary distribution correction estimation." ICML 2021.

[2] Mao, Liyuan, et al. "Odice: Revealing the mystery of distribution correction estimation via orthogonal-gradient update." ICLR 2024.

[3] Nachum, Ofir, et al. "Algaedice: Policy gradient from arbitrary experience." arXiv preprint arXiv:1912.02074 (2019).

[4] Peng, Xue Bin, et al. "Advantage-weighted regression: Simple and scalable off-policy reinforcement learning." arXiv preprint arXiv:1910.00177 (2019).

[5] Xu, Haoran, et al. "Offline rl with no ood actions: In-sample learning via implicit value regularization." ICLR 2023.
------Thank you for continuing to read!------

**Q3: How sensitive is the performance to the accuracy of the model?**

Thanks for your question. To empirically investigate SCAS under different dynamics model errors, we run SCAS using different checkpoints of the trained dynamics model, which are obtained at different steps in the dynamics model training process. The model error is controlled by the number of trained steps. The results are shown in **Figure 1** of the PDF. The figure plots the training loss of the dynamics model $M_\omega$ and the corresponding normalized return of SCAS over 5 random seeds. We observe that the performance of SCAS increases with the number of trained steps of the dynamics model (i.e. the accuracy of the model) and stabilizes at a high level.

**Q4: Is it feasible to apply the regularizer across different offline RL objectives, such as CQL, IQL, and others?**

Yes, the SCAS regularizer is compatible with various offline RL methods. We have conducted experiments to combine SCAS with CQL, IQL, and TD3BC. Comparisons between these combined algorithms and the original ones are shown in **Table 2** of the PDF. We find that applying the SCAS regularizer leads to improved performance for all these popular algorithms, which could be attributed to the OOD state correction effects of SCAS. However, we also find that these combined methods do not achieve better performance than the original SCAS (comparable on most tasks and worse on some tasks). We hypothesize that this is because SCAS already has the effect of OOD action suppression, and when combined with offline RL objectives that also aim for OOD action suppression, it may become overly conservative. As a result, the combined algorithms may perform worse than the original SCAS on some sub-optimal datasets.

**Q5: Could this approach be extended to also incorporate the model during testing time?**

This is a good perspective! Indeed, SCAS utilizes the trained model as a regularizer during policy training. Incorporating the model during testing time for OOD state correction can be an alternative approach (or an enhancement to SCAS) to achieving test-time OOD state correction. These two directions appear to be completely different in algorithmic designs, and it would be an interesting direction for future work. Thanks for the insightful comment.

**Reference**

[1] Hong, Zhang-Wei, et al. "Beyond uniform sampling: Offline reinforcement learning with imbalanced datasets." NeurIPS 2023.

[2] Yang, Rui, et al. "Rorl: Robust offline reinforcement learning via conservative smoothing." NeurIPS 2022.

[3] An, Gaon, et al. "Uncertainty-based offline reinforcement learning with diversified q-ensemble." NeurIPS 2021.

[4] Xu, Haoran, et al. "Offline rl with no ood actions: In-sample learning via implicit value regularization." ICLR 2023.
------Thank you for continuing to read!------

**Q6: In Figure 2, why do Off-policy, SDC, and OSR only have one dot? Shouldn't it also be a line?**

We apologize for the confusion. Because the Q values of these methods diverge in the early learning stage, showing the complete lines would result in an excessive range on the vertical axis. For clearer presentation, we have replotted this figure in **Figure 2** of the PDF, which adds more data points and shows the lines in the early learning stage.

**Q7: How are the results in Table 1 recorded? Final round results, last 5 epoch average, or another method?**

Our evaluation criteria follow those used in most previous works. We evaluated the last trained policy, with each experiment repeated over 5 random seeds. For the Gym locomotion tasks, we average returns over 10 evaluation trajectories, while for the Antmaze tasks, we average over 100 evaluation trajectories.

**Q8: Where do the baseline method results come from?**

The baseline method results are obtained as follows. We re-run OSR on all datasets using their official codebase and tune the hyperparameters for each dataset as specified in their paper. We implement SDC and re-run it on all datasets. We tune the SDC-related hyperparameters as specified in their paper, and sweep the CQL-related hyperparameters in {1,2,5,10,20} for each dataset. We re-run OneStep RL on all datasets using its official codebase and the default hyperparameters. We implement BC based on the TD3+BC repository and re-run it on all datasets. The results of other baselines are taken from the PBRL paper[3] and SPOT paper[2].

**Reference**

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

[2] Wu et al. Supported policy optimization for offline reinforcement learning. NeurIPS 2022.

[3] Bai et al. Pessimistic Bootstrapping for Uncertainty-Driven Offline Reinforcement Learning. ICLR 2022.