SlowGuess/NeurIPS_2023_OCR huggingface dataset:
where  $\phi$  and  $\psi$  are neural networks. Next, we present how to train both neural network models. As the dataset often has limited coverage on state-action space, it is preferable to add a KL-divergence regularization  $D_{KL}(\mathcal{D}_w||\mathcal{D})$  to the objective in Equation 10, as proposed in Zhan et al. [43]. This regularization keeps the state-action distribution  $\mathcal{D}_w$  induced by the learned weighting  $w$  close to the original dataset  $\mathcal{D}$ , preventing  $w_{\phi,\psi}(s,a)$  from overfitting to a few rare state-action pairs in  $\mathcal{D}$ . Note that this does not prevent the learned weightings to provide a better data distribution for regularized offline RL algorithms. See Appendix A.2 for the detailed discussion. Another technical difficulty on training  $w$  is that it is difficult to impose Bellman flow conservation constraint in Equation 10 at every state in the state space since only limited coverage of states are available in the dataset. Thus, we instead use penalty method [3] to penalize the solution of  $w_{\phi,\psi}$  on violating this constraint in expectation. As a result, we optimize  $w_{\phi,\psi}$  for Equation 10 using stochastic gradient ascent to optimize the following objective (details can be found in Appendix A.3):
$$
\max  _ {\phi , \psi} \mathbb {E} _ {(s, a, s ^ {\prime}) \sim \mathcal {D}} \left[ \underbrace {w _ {\phi , \psi} (s , a) r (s , a)} _ {\text {R e t u r n}} - \lambda_ {F} \underbrace {\left(w _ {\phi} \left(s ^ {\prime}\right) - w _ {\phi , \psi} (s , a)\right) ^ {2}} _ {\text {B e l l m a n f l o w c o n s e r v a t i o n p e n a l t y}} \right] - \lambda_ {K} \underbrace {D _ {K L} (\mathcal {D} _ {w} | | \mathcal {D})} _ {\text {K L r e g u l a r i z a t i o n}}, \tag {13}
$$
where  $s'$  denotes the next state observed after taking action  $a$  at state  $s$ , and  $\lambda_F, \lambda_K \in \mathbb{R}^+$  denote the strength of both penalty terms. Note that the goal of our work is not to propose a new off-policy evaluation method, but to motivate ours in the specific objective to optimize the importance weighting for training offline RL algorithms. Importantly, our approach differs from previous off-policy evaluation methods [32, 26, 44], as further discussed in the related works (Section 6).
Applying the weighting to offline RL. The weighing function  $w_{\phi, \psi}$  could be pre-trained before training the policy, but this would introduce another hyperparameter: the number of pretraining iterations. As a consequence, we opt to train  $w_{\phi, \psi}$  in parallel with the offline RL algorithm (i.e., value functions and policy). In our experiments, we perform one iteration of offline RL update pairs with one iteration of weighting function update. We also found that weighting both  $\widehat{J}_{\mathcal{D}}^{\gamma}(\pi)$  and  $\mathcal{C}(s, a)$  at each state-action pairs sampled from the dataset  $\mathcal{D}$  with  $w_{\phi, \psi}(s, a)$  performs better than solely weighting the regularization term  $\mathcal{C}(s, a)$ . For example, when weighting the training objective of implicit Q-learning (IQL) [18] (an offline RL method), the weighted objective  $J_{\mathcal{D}_w}(\pi)$  is:  $\mathbb{E}_{(s, a) \sim \mathcal{D}}[w_{\phi, \psi}(s, a)A(s, a)\log \pi(a|s)]$ , where  $A(s, a)$  denotes advantage values. Please, see Appendix A.3 for implementation details. We hypothesize that weighting both the policy optimization objective  $\widehat{J}_{\mathcal{D}}^{\gamma}(\pi)$  and regularization  $\mathcal{C}(s, a)$  in the same distribution (i.e., same importance weights) is needed to prevent policy  $\pi$  increasing  $\widehat{J}_{\mathcal{D}}^{\gamma}(\pi)$  by exploiting out-of-distribution actions on states with lower weights  $w_{\psi, \phi}(s, a)$ , which could lead to poor performance [8]. Appendix A.5.5 compares weighting both and only one objective. The training procedure is outlined in Algorithm 1.
Our experiments aim to answer whether our density-ratio weighting (DW) (Section 4) approach can improve the performance of offline RL algorithms with different types of imbalanced datasets (Section 3). Prior work on imbalanced datasets [12] focused exclusively on imbalanced datasets with trajectories originating from a similar initial state. However, in real-world scenarios, trajectories can be collected from diverse initial states. For instance, when collecting datasets for self-driving cars, it is likely that drivers initiate the recording of trajectories from drastically different initial locations. We
Algorithm 1 Density-ratio weighting with generic offline RL algorithms (details in Appendix A.3)
1: Input: Dataset  $\mathcal{D}$
2: Initialize policy  $\pi$  and weighting function  $w_{\phi, \psi}$
3: while not converged do
4: Sample a batch  $\mathcal{B}$  of tuples of states, actions, rewards, and next states  $(s,a,r,s^{\prime})$  from  $\mathcal{D}$
5: Update  $w_{\phi, \psi}$  with batch  $\mathcal{B}$  using Equation 13
6: Update policy  $\pi$  and value function with an offline RL training objective with weights  $w_{\phi, \psi}(s, a)$  and  $\mathcal{B}$  (e.g., [22, 7, 18])
7: end while
found that imbalanced datasets with diverse initial states exhibit a long-tailed distribution of trajectory returns, while those with similar initial states show a bimodal distribution (see Appendix A.4 for details). As diversity of initial states affects the type of imbalance, we focus our experimentation on the two types of datasets: (i) Trajectories with similar initial states and (ii) Trajectories with diverse initial states.
Following the protocol in prior offline RL benchmarking [6], we develop representative datasets of each type using the locomotion tasks from the D4RL Gym suite. Our datasets are generated by combining  $1 - \sigma\%$  of trajectories from the random-v2 dataset (low-performing) and  $\sigma\%$  of trajectories from the medium-v2 or expert-v2 dataset (high-performing) for each locomotion environment in the D4RL benchmark. For instance, a dataset that combines  $1 - \sigma\%$  of random and  $\sigma\%$  of medium trajectories is denoted as random-medium- $\sigma\%$ . We evaluate our method and the baselines on these imbalanced datasets across four  $\sigma \in \{1, 5, 10, 50\}$ , four environments. Both types of datasets are briefly illustrated below and detailed in Appendix A.4. Additionally, we present the results on the rest of original D4RL datasets in Appendix A.5.
(i) Trajectories with similar initial states. This type of datasets was proposed in [12], mixing trajectories gathered by high- and low-performing policies, as described in Section 3. Each trajectory is collected by rolling out a policy starting from similar initial states until reaching timelimit or terminal states. We consider a variant of smaller versions of these datasets that have small number of trajectories, where each dataset contains 50,000 state-action pairs, which is 20 times smaller. These smaller datasets can test if a method overfits to small amounts of data from high-performing policies.
(ii) Trajectories with diverse initial states. Trajectories in this type of dataset start from a wider range of initial states and have varying lengths. One real-world example of this type of dataset is a collection of driving behaviors obtained from a fleet of self-driving cars. The dataset might encompass partial trajectories capturing diverse driving behaviors, although not every trajectory accomplishes the desired driving task of going from one specific location to the other. As not all kinds of driving behaviors occur with equal frequency, such a dataset is likely to be imbalanced, with certain driving behaviors being underrepresented
5.1 Evaluation Setup
Baselines and prior methods. We consider uniform sampling (denoted as Uniform) as the primary baseline for comparison. In addition, we compare our method with two existing approaches for improving offline RL performance on imbalanced datasets: advantage-weighting (AW), proposed in the recent work by [12] and percentage-filtering (PF) [5]. Both AW and PF sample state-action pairs with probabilities determined by the trajectory's return to which they belong. The sampling probabilities for AW and PF are given as follows:
$$
\mathcal {P} _ {\mathrm {A W}} \left(s _ {t} ^ {i}, a _ {t} ^ {i}\right) \propto \exp \left(\left(G \left(\tau_ {i}\right) - V _ {0} \left(s _ {0} ^ {i}\right)\right) / \eta\right) \quad (\text {A d v a n t a g e - w e i g h t i n g}) \tag {14}
$$
$$
\mathcal {P} _ {\mathrm {PF}} \left(s _ {t} ^ {i}, a _ {t} ^ {i}\right) \propto \mathbb {1} \left[ G \left(\tau_ {i}\right) \geq G _ {K \%} \right] \tag{15}
$$
where  $(s_t^i, a_t^i)$  denotes the state-action pair at timestep  $t$  of trajectory  $\tau_i$ .  $G_{K\%}$  represents a threshold for selecting the top- $K\%$  of trajectories, with  $K$  chosen from  $\{10, 20, 50\}$  as practiced in [12].  $V(s_0^i)$  denotes the value of the initial state  $s_0^i$  in trajectory  $\tau_i$ , and the coefficient  $\eta$  represents the temperature coefficient in a Boltzmann distribution. We consider three levels of  $\eta$ : low (L), medium (M), and high (H) in our experiments. Further details of the hyperparameter setup can be found in Appendix A.4. For the following experiments, we implement our DW and the above baselines on the top of state-of-the-art offline RL algorithms: Conservative Q-Learning (CQL) [22], Implicit
Q-Learning (IQL) [19], and TD3BC [7]. Note that as AW can provide a better initial sampling distribution to train the weighting function in DW, we initialize training DW with AW sampling (denoted as  $DW$ -AW) and initialize training DW with uniform sampling (denoted as  $DW$ -Uniform) in the following experiments. We refer the readers to Appendix A.3 for the implementation details.
Evaluation metrics. Following the settings of [6], we train all algorithms for one million gradient steps with three random seeds in each dataset. We evaluate the performance of policies acquired through each method in the environment corresponding to the dataset by conducting 20 episodes every 1000 gradient step. To determine the policy's performance at a given random seed, we compute the average returns over 20 episodes during the final 10 evaluation rounds, each round separated by 1000 gradient steps. We chose to average over the performance at the last 10 rounds rather than solely at the last evaluation round because we observed that the performance of offline RL algorithms oscillates during gradient steps. The main performance metric reported is the interquartile mean (IQM) [1] of the normalized performance across multiple datasets, along with its  $95\%$  confidence interval calculated using the bootstrapping method. As suggested in [1], IQM is a robust measure of central tendency by discarding the top and bottom  $25\%$  of samples, making it less sensitive to outliers.
5.2 Scenario (i): Trajectories with Similar Initial States
Figure 3a shows IQM of the normalized return for thirty-two different datasets where trajectories start from similar initial states (Section 5). For all the datasets, we use the best hyperparameters for AW and PF found in [12] and the hyperparameters for DW-AW and DW-Uniform are presented in Appendix A.4. The results demonstrate that both DW-AW and DW-Uniform outperform the uniform sampling approach confirming the effectiveness of our method. Moreover, combining DW with AW enhances the performance of DW, indicating that DW can benefit from the advantages of AW. This is likely because AW can provide a good initial sampling distribution to start training the weighting function in DW. While DW-Uniform did not exceed the performance of AW in this experiment, it should be noted that our method can be applied when datasets are not curated with trajectories such as reset-free or play style datasets where data is not collected in an episodic manner. This is useful for continuing tasks, where an agent (data curator) performs a task infinitely without termination (i.e., locomotion).
Limited size datasets. Figure 3b presents the results on smaller versions of 8 of these datasets used in Figure 3a. Note that as we observe the higher temperature  $\eta$  enables AW with CQL to perform better in this type of dataset, we additionally consider AW-XH (extra high temperature) for comparison to provide AW with as fair a comparison point as possible. Further details on the hyperparameter settings can be found in Appendix A.4. Our methods consistently achieve significantly higher returns compared to AW and PF when combined with CQL. This suggests that our methods effectively utilize scarce data in smaller datasets better than weighted sampling approaches (AW and PF) that rely on episodic trajectory-based returns rather than purely transition level optimization like DW. In the case of IQL and TD3BC, we see a clear performance improvement of our methods over uniform sampling and PF while our methods perform on par with AW. For IQL, we hypothesize that this is because IQL is less prone to overfitting on small amounts of data due to its weighted behavior cloning objective [18], which always uses the in-distribution actions. However, it is worth noting that IQL falls short in performance compared to CQL with DW-Uniform. This suggests that IQL may primarily focus on replicating behaviors from high-return trajectories instead of surpassing them, as it lacks the explicit dynamic programming used in CQL.
Takeaway. Since CQL with DW-AW outperforms the other two offline RL algorithms in both dataset types, our suggestion is to opt for CQL with DW-AW, especially when dealing with datasets that might exhibit an imbalance and include trajectories originating from comparable initial states.
5.3 Scenario (ii): Trajectories with Diverse Initial States
Figure 4 presents the results on thirty-two datasets of trajectories with diverse initial states (Section 5). We observe that uniform sampling's performance drops significantly in these datasets compared to trajectories with similar initial states, indicating that the presence of diverse initial states exacerbates the impact of imbalance. Both of our methods consistently outperform all other approaches considered in Section 5.2, including AW and PF methods. Notably, even the best-performing variant of AW
(a) Results on imbalanced datasets of trajectories with similar initial states (Section 5.1).
(b) Results on smaller version of datasets used in Figure 3a.
Figure 3: (a) Our methods, DW-AW and DW-Uniform, achieve higher return than Uniform, indicating that DW can enhance the performance of offline RL algorithms on imbalanced datasets. Note that our methods in IQL, although not surpassing AW and PF-10% in performance, ours can be applied to offline RL dataset that are not curated with trajectories. (b) Our methods outperform Uniform in CQL, IQL, and TD3BC, indicating no significant overfitting in smaller datasets. DW-AW demonstrates superior returns compared to AW and PF, particularly in CQL, indicating our method effectively leverages limited data. IQL shows limited gains likely due to its difficulties in utilizing data from the rest of low-return trajectories in the dataset (see Section 5.2).
(AW-M) falls short of matching the performance of our DW-AW, demonstrating the effectiveness of DW in leveraging the initial sampling distribution provided by AW and furthering its performance. The performance degradation of AW can be attributed to the presence of diverse initial states and varying trajectory lengths in these datasets. In such cases, state-action pairs in trajectories with high returns are not necessarily generated by high-performing policies. For instance, a sub-optimal policy can also easily reach the goal and achieve a high return if it starts close to the goal (i.e., lucky initial states). Consequently, over-sampling state-action pairs from high-return trajectories can introduce bias towards data in trajectories starting from lucky initial states. Although AW attempts to address this issue by subtracting the expected return of initial states (see Section 5.1), our results show that AW has limited success in addressing this issue. This is because the estimated expected returns of initial states can be inaccurate since AW uses the trajectories' returns in the dataset to estimate initial states' returns (i.e., Monte Carlo estimates). The trajectories in the dataset are finite in length, which makes the Monte Carlo estimates of expected returns inaccurate. To conclude, when an imbalanced dataset consists of trajectories starting from diverse initial states, we recommend using DW-AW to re-weight the training objectives in offline RL algorithms.
Our approach builds upon recent advances in off-policy evaluation techniques, specifically density-ratio importance correction estimation (DiCE) [32]. DiCE has been primarily used for policy evaluation [29, 32, 10], while our method makes use of DiCE (i.e., the learned importance weights) to re-weight samples for offline RL algorithms. Recent works [43, 37, 33, 26] optimize the policy using DiCE via re-weighting behavior cloning with DiCE, while we found it fails to match offline RL algorithms' performance even in datasets with plenty of expert demonstration (Appendix A.5).
Offline imitation learning approaches [15, 30, 41] also consider imbalanced datasets similar to ours. However, these methods assume prior knowledge of which data points are generated by experts,
Figure 4: Results on imbalanced datasets with trajectories starting from diverse initial states (Section 5.3). Compared to Figure 3a, the performance of uniform sampling and AW decrease, showing that diverse initial states exacerbate the issue of imbalance. Our methods, DW-AW and DW-Uniform, achieve higher return than all the baselines, which suggests DW is advantageous in broader types of imbalanced datasets.
while our approach does not rely on such information. Furthermore, our method can effectively handle datasets that include a mixture of medium-level policies and low-performing policies, whereas existing approaches often rely on expert-labeled data.
Multi-task offline RL algorithms [42, 14] filter data relevant to the current task of interest from datasets collected from multiple task. For example, Yu et al. [42] employ task relevance estimation based on Q-value differences between tasks. While our motivation aligns with data filtering, our problem setting differs as we do not assume knowledge of task identifiers associated with the data points. Additionally, our dataset comprises varying levels of performance within the same task, while existing works mix data from different tasks.
Support constraints [20, 38, 2, 40] have been proposed as an alternative approach to prevent offline RL algorithms from exploiting out-of-distribution actions, distinct from distributional constraints used in state-of-the-art methods [22, 7, 18]. While support constraints theoretically suit imbalanced data, the prior work [38] found that support constraints have not shown significant improvements beyond distributional constraint-based algorithms. Note that our method is independent of the constraint used in offline RL algorithms. Thus support constraints is orthogonal to our approach.
7 Conclusion, Future Directions, and Limitations
Our method, density-ratio weighting (DW) improves the performance of state-of-the-art offline RL algorithms [22, 18] over 72 imbalanced datasets with varying difficulties. In particular, our method exhibits substantial improvements in more challenging and practical datasets where the trajectories in the dataset start from diverse initial states and only limited amount of data are available. Future works can explore other optimization techniques to better address the Bellman flow conservation constraint in importance weights optimization (e.g., Augmented Lagrangian method [37]). Additionally, it would be valuable to study the impact of violating this constraint on the effectiveness of importance-weighted offline RL algorithms.
Limitations. Although our method improves performance by optimizing sample weights, we lack theoretical guarantees due to the absence of a unified theoretical analysis on the dependence of state-of-the-art offline RL algorithms on imbalanced data distribution. While some theoretical works [4, 36] have analyzed the interplay between data distribution and offline RL algorithm performance, they primarily focus on specific algorithms that differ significantly from the practical state-of-the-art offline RL algorithms.
Author Contributions
- Zhang-Wei Hong: Led the project and the writing of the paper, implemented the method, and conducted the experiments.
- Aviral Kumar: Advised the project in terms of theory, algorithm development, and experiment design. Revised the paper and positioned the paper in the field.
- Sathwik Karnik: Prepared the datasets and proofread the paper.
- Abhishek Bhandwaldar: Helpd scaling up experiments in the cluster.
- Akash Srivastava: Advised the project in the details of the practical and theoretical algorithm design and coordinated the compute.
- Joni Pajarinen: Advised the project in the details of the practical and theoretical algorithm design and experiment designs.
- Romain Laroche: Advised the project in the theory of the algorithms and dataset designs.
- Abhishek Gupta: Advised the project in terms of theory, algorithm development, and experiment design. Revised the paper and positioned the paper in the field.
- Pulkit Agrawal: Coordinated the project, revised the paper, and positioned the paper in the field.
We thank members of the Improbable AI Lab for helpful discussions and feedback. We are grateful to MIT Supercloud and the Lincoln Laboratory Supercomputing Center for providing HPC resources. This research was supported in part by the MIT-IBM Watson AI Lab, an AWS MLRA research grant, Google cloud credits provided as part of Google-MIT support, DARPA Machine Common Sense Program, ARO MURI under Grant Number W911NF-21-1-0328, ONR MURI under Grant Number N00014-22-1-2740, and by the United States Air Force Artificial Intelligence Accelerator under Cooperative Agreement Number FA8750-19-2-1000. The views and conclusions contained in this document are those of the authors and should not be interpreted as representing the official policies, either expressed or implied, of the Army Research Office or the United States Air Force or the U.S. Government. The U.S. Government is authorized to reproduce and distribute reprints for Government purposes notwithstanding any copyright notation herein.
[1] R. Agarwal, M. Schwarzer, P. S. Castro, A. C. Courville, and M. Bellemare. Deep reinforcement learning at the edge of the statistical precipice. Advances in neural information processing systems, 34:29304-29320, 2021.
[2] G. An, S. Moon, J.-H. Kim, and H. O. Song. Uncertainty-based offline reinforcement learning with diversified q-ensemble. Advances in neural information processing systems, 34:7436-7447, 2021.
[3] S. P. Boyd and L. Vandenberghe. Convex optimization. Cambridge university press, 2004.
[4] J. Buckman, C. Gelada, and M. G. Bellemare. The importance of pessimism in fixed-dataset policy optimization. arXiv preprint arXiv:2009.06799, 2020.
[5] L. Chen, K. Lu, A. Rajeswaran, K. Lee, A. Grover, M. Laskin, P. Abbeel, A. Srinivas, and I. Mordatch. Decision transformer: Reinforcement learning via sequence modeling. arXiv preprint arXiv:2106.01345, 2021.
[6] J. Fu, A. Kumar, O. Nachum, G. Tucker, and S. Levine. D4rl: Datasets for deep data-driven reinforcement learning. arXiv preprint arXiv:2004.07219, 2020.
[7] S. Fujimoto and S. S. Gu. A minimalist approach to offline reinforcement learning. Advances in neural information processing systems, 34:20132-20145, 2021.
[8] S. Fujimoto, D. Meger, and D. Precup. Off-policy deep reinforcement learning without exploration. arXiv preprint arXiv:1812.02900, 2018.
[9] S. Fujimoto, H. van Hoof, and D. Meger. Addressing function approximation error in actor-critic methods. arXiv preprint arXiv:1802.09477, 2018.