Beyond Uniform Sampling: Offline Reinforcement Learning with Imbalanced Datasets

Summary

This paper investigates the offline RL with the imbalanced dataset setting. The proposed method, which optimizes a parameterized density-ratio weighting (DW) model for each transition, anchors a policy close to the good part of trajectories in the dataset. The proposed method is shown to outperform state-of-the-art methods on 72 imbalanced datasets both in varying imbalance and initial state distribution.

Strengths

1. Overall good writing quality and clarity. 2. The proposed approach is well-motivated by adopting existing approaches. 3. Sufficient technical details for reproducing the experiments.

Weaknesses

1. Novelty limited: Fixed conservativeness is defective for imbalanced (heterogeneous) datasets and reweighting or filtering samples/trajectories are not new in offline RL [1, 2, 3]. On the other hand, this paper does not connect with the previous most similar work [1] and explains why DW is better than AW, though it claims that AW is empirically prone to overfitting. 2. Lack of details for the dataset setup. 3. Lack of comparison: the imbalanced dataset is also called the noisy dataset [4]. If the authors follow the setup of the dataset from AW [1], the same as the noisy data setup [4], it is reasonable to include more comparisons between SQL/EQL enforcing a Q-function weighting on the transitions and show the edge of DW. 3. Experiments between AW and DW: In Figure 2, AW-M is better than DW-Uniform in three of four tasks. However, the authors claim that they are better than AW. More explanation is needed about why DW is better than AW.

Questions

4. It is reasonable to set $\gamma=1$ to avoid the estimation of $\rho_0$. Have the authors analyzed the role of the discount factor for offline RL [5]? It will be appreciated if any theoretical analysis or/and experiments are provided. 5. Lacking the AW-L benchmark in Figure 2(b). The medium temp is extremely beneficial in IQL (see Figure 2(b)), why do not present the performance of AW-L? 7. In Figure 3, It is counterintuitive that DW is not influenced by diverse initial states and varying trajectories because DW does not consider the initial state distribution. Could the authors explain this? 8. Thank the authors for providing technical details including their hyperparameter search range with $\lambda_f$ and $\lambda_k$. However, I am a bit concerned that the necessity of penalties, i.e., Bellman flow and KL-divergence. Could the authors please introduce more experiences to show their necessities and the guidance for hyperparameter research with new tasks? [1] Hong Z W, Agrawal P, des Combes R T, et al. Harnessing Mixed Offline Reinforcement Learning Datasets via Trajectory Weighting[C]//The Eleventh International Conference on Learning Representations. 2023. [2] Brandfonbrener D, Whitney W F, Ranganath R, et al. Quantile filtered imitation learning[J]. arXiv preprint arXiv:2112.00950, 2021. [3] Wu Y, Zhai S, Srivastava N, et al. Uncertainty weighted actor-critic for offline reinforcement learning[J]. arXiv preprint arXiv:2105.08140, 2021. [4] Xu H, Jiang L, Li J, et al. Offline rl with no ood actions: In-sample learning via implicit value regularization[J]. arXiv preprint arXiv:2303.15810, 2023. [5] Hu H, Yang Y, Zhao Q, et al. On the role of discount factor in offline reinforcement learning[C]//International Conference on Machine Learning. PMLR, 2022: 9072-9098.

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

4 excellent

Contribution

3 good

Limitations

See questions and weaknesses. --- raise the score to 6 after reviewing the response from the authors. --- **Minors**: 1. Cite twice: The classic results in [34] indicate that a policy’s expected return can be expressed in terms of its stationary state-action distribution [34]. 2. IQL training object for policy improvement (line 234) is wrong. 3. What is DW-uniform? Only use DW? I suggest changing the name of DW-uniform because *uniform* will cause misunderstanding that DW-uniform is a combination sampling method between DW and Uniform.

Summary

Authors propose a new method for weighting the samples from datasets in order to mimic the sampling from dataset which is collected by the better policy the the behavioral policy. This method can be integrated into any algorithm and combined with previous approaches.

Strengths

Method seems to be not that hard to implement and it shouldn't introduce a notable computational overhead while can be applied with the arbitary algorithm.On imballanced locomotion datasets generated by the authors proposed approach helps to achieve much better results than other approaches. Benchmarking over high number of the datasets is another strength.

Weaknesses

While approach helps to improve performance when D4RL locomotion datasets are mixed with the random datasets, it seems like there is no advantage when applied to the original mixed datasets (medium-replay, medium-expert, full-replay) or the performance even drops on those if we look at Tables 4 and 5.

Questions

What are the avereged scores across domains in tables 4 and 5? As mentioned in "Weaknesses" it seems like the proposed method does not benefit a lot (if benefit at all) when applied to the original D4RL datasets. Could you please add those numbers to see what is the situation on average? I mean average scores over locomotion, antmaze, kitchen, adroit, generated datasets and averaged scores over all of the domains. How does the method affect algorithms time needed for training?

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

3 good

Contribution

2 fair

Limitations

Beside the uncetainty whether approach helps when dataset is not imbalanced that hard there is another limitation in design. If I understand approach correctly, it can't be applied to the datasets which require trajectories stiching to complete task. The example of such task is AntMaze and authors' approach mostly decrease performance when applied to those datasets.

Summary

This paper proposes a method to improve offline reinforcement learning (RL) performance on imbalanced datasets, where most of the data comes from low-performing policies and only a few from high-performing ones. The method, called density-ratio weighting (DW), optimizes the importance sampling weights to emulate sampling data from a data distribution generated by a nearly optimal policy. The paper shows that DW can enhance the performance of state-of-the-art offline RL algorithms on 72 imbalanced datasets with varying types of imbalance.

Strengths

- It proposes a novel method of DW that optimizes the importance sampling weights to emulate sampling data from a data distribution generated by a nearly optimal policy, rather than the original dataset. - It demonstrates overall performance gains over SoTA offline RL algorithms and other baselines on imbalanced datasets with varying types and degrees of imbalance. - Its written style is clear and easy to follow.

Weaknesses

- The theoretical analysis is relatively weak. - Slightly lack of discussion with a type of offline RL algorithms which filter out low-performance trajectories.

Questions

- Though obtaining a performance bound is quite difficult as describing the imbalanced dataset is hard, is it possible to prove the optimization with the proposed loss function will converge? For example, I guess, derive an iteration equation of $w$ as the Bellman equation, and prove the operator is a contraction mapping. - In Section 6 related work, the paper discusses offline imitation learning approaches. It claims that these approaches "assume prior knowledge of which data points are generated by experts". However, there are some offline imitation or RL methods which only use the reward signals in the dataset and filter out low-quality data to improve the final performance, such as BAIL [1] and COIL [2]. I am curious about what the advantages of DW are compared to them, especially for BAIL which does not need trajectory information. Also, it will be better if you can try BAIL on your imbalanced dataset and show the results. [1] Chen, Xinyue, et al. "Bail: Best-action imitation learning for batch deep reinforcement learning." Advances in Neural Information Processing Systems 33 (2020): 18353-18363. [2] Liu, Minghuan, et al. "Curriculum offline imitating learning." Advances in Neural Information Processing Systems 34 (2021): 6266-6277.

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

Yes.

Summary

The authors consider an offline RL problem where the offline dataset has a small number of high-reward episodes and a larger pool of low-reward episodes. They argue this setting is fairly common in reality, since generating high-reward episodes is often higher effort. Considering offline RL objective functions that are structured like "maximize reward while staying close to the behavior policy", the argument is that staying close to the behavior policy causes optimization to be too conservative / biased too much towards the large pool of low-reward episodes. Borrowing the DiCE techniques from off-policy evaluation, the proposed method is to learn importance weights to bias the "stay close to data" objective towards the high-reward (expert) trajectories over the support of the original offline dataset. The learning of these weights is done standard to the DiCE style techniques, although in practice the importance weights are given a KL constraint to stay close to the original dataset to avoid collapsing the distribution to a few rare examples, and the importance weights are also applied to the reward-maximizing terms in offline RL (even though theoretically this should not be required). This is then compared with two other approaches that bias the offline distribution towards higher return episodes - sampling based on the advantage of the offline episode, or sampling the top K% of the offline data. The best results are found by combining advantage weighting with density ratio weighting

Strengths

I have some objections about claiming that offline RL methods have not considered the downsides of constraining to poor return trajectories. (More on this later). But in general, I agree that many datasets often have a smaller number of good examples, importance weights to adjust the conservatism penalty make sense, and using methods from the OPE literature is a reasonable way to learn those importance weights. The evaluation is also extensive enough for me to trust the results.

Weaknesses

There have been methods that attempt to constraint the policy to the support of the offline dataset, rather than the behavior of the offline dataset. BEAR and BRAC come to mind due to using Kernel MMD to measure policy divergence (and they ought to be cited.) In practice these methods have underperformed CQL and TD3+BC so I think it is okay to benchmark against just CQL, but it is still a notable omission. I also suspect that the KL penalty is one of the more important hyperparams, would need to be tuned separately per dataset (as the acceptable level of deviation will change depending on dataset composition), and there is not much guidance on how to set this parameter. This weakness is common to many offline RL methods though (i.e. the weight of CQL penalty also needs to be tuned).

Questions

Could the authors comment on how the importance weights $w(s,a)$ evolve over time? To me the most dangerous outcome is that $w(s,a)$ learning becomes unstable, or training overfits too heavily due to upsampling a smaller section of the data. Is there any way to compare how the weighting of $w(s,a)$ compares to methods like top K%, in how much they consider different examples in the data?

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

4 excellent

Presentation

3 good

Contribution

3 good

Limitations

Seems fine.

Response to the Authors' Rebuttal by Reviewer XiFQ

I have carefully reviewed the comments from other reviewers, and the authors have addressed most of my concerns. The additional experiments provided sufficiently demonstrate the advantages of density ratio optimization. However, I still have one minor concern. Regarding my previous concerns about Q1 and Q3, the authors chose to set $\gamma=1$ to bypass the dependence on the initial state distribution, i.e., $\rho_0$. Nonetheless, I believe that setting $\gamma=1$, which is agnostic with the initial state distribution, comes with sacrifice. Could the authors discuss the limitations associated with this choice? Overall, I appreciate the detailed response and have decided to raise my rating to 6.