Offline Reinforcement Learning with OOD State Correction and OOD Action Suppression

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.

Strengths

The authors highlighted the significance of OOD state correction in offline RL, a problem other researchers often neglect. Also, multiple experiments were conducted to show the competence of SCAS and its robustness towards the choice of hyperparameters.

Weaknesses

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

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.

Rating

6

Confidence

3

Soundness

4

Presentation

2

Contribution

3

Limitations

The authors have adequately addressed their work's limitations and potential negative societal impact.

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.

Strengths

**1. Theoretical connection between OOD state correction and OOD action suppression** A notable strength of this paper is the establishment of a theoretical link between OOD state correction and OOD action suppression. The authors articulate how the mechanism designed to correct OOD states implicitly suppresses OOD actions, by showing an optimal policy of the regularized objective will produce actions inside of the support of behavior policy for every dataset state. This implicit suppression also explains successful empirical results on offline RL benchmarks. **2. Robustness to Hyperparameter Selection** As detailed in Section 6.2 and illustrated in Figure 4, the proposed approach demonstrates considerable robustness to variations in hyperparameter settings. This attribute is particularly valuable in practical applications of offline RL where optimal hyperparameter settings can be elusive or computationally expensive to determine. **3. Evaluation on standard benchmarks and implementation code is attached for reproduction.** The paper's evaluation methodology is another major strength. The authors have rigorously tested their approach on standard benchmarks, D4RL and NeoRL, providing a comprehensive assessment of its performance relative to existing methods. Moreover, the inclusion of implementation code enhances the paper’s contribution by facilitating reproducibility and further experimentation.

Weaknesses

**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

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?

Rating

6

Confidence

4

Soundness

3

Presentation

2

Contribution

2

Limitations

The authors included their limitations in Section 7 and limitations are addressed appropriately.

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.

Strengths

1. The paper presents an intuitive solution to the OOD state iusse, while also mitigating the overestimation of OOD actions. 2. The algorithm is highly robust, requiring only a single set of hyperparameters across different environments. This robustness is particularly important in the context of offline reinforcement learning, as online interaction for tuning parameters is costly or even dangerous.

Weaknesses

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

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.

Rating

6

Confidence

3

Soundness

2

Presentation

3

Contribution

3

Limitations

Please see weakness.

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.

Strengths

I found it an interesting topic to consider the OOD state issue. Most existing offline RL methods focus on addressing the OOD action issue, while this paper highlights that OOD states may also cause training collapse. Another good property of the method is that it only needs one single hyperparameter configuration to achieve good empirical performance. The theory is also solid.

Weaknesses

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

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?

Rating

6

Confidence

4

Soundness

3

Presentation

3

Contribution

3

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.