Title: Beyond Uniform Sampling: Offline Reinforcement Learning with Imbalanced Datasets
Venue: NeurIPS 2023
Accepted: yes
Keywords: offline reinforcement learning

reinforcement learning

sampling

experience replay
Abstract: Offline reinforcement learning (RL) enables learning a decision-making policy without interaction with the environment. This makes it particularly beneficial in situations where such interactions are costly. 
However, a known challenge for offline RL algorithms is the distributional mismatch between the state-action distributions of the learned policy and the dataset, which can significantly impact performance. State-of-the-art algorithms address it by constraining the policy to align with the state-action pairs in the dataset. However, this strategy struggles on datasets that predominantly consist of trajectories collected by low-performing policies and only a few trajectories from high-performing ones. Indeed, the constraint to align with the data leads the policy to imitate low-performing behaviors predominating the dataset. Our key insight to address this issue is to constrain the policy to the policy that collected the good parts of the dataset rather than all data. To this end, we optimize the importance sampling weights to emulate sampling data from a data distribution generated by a nearly optimal policy. Our method exhibits considerable performance gains (up to five times better) over the existing approaches in state-of-the-art offline RL algorithms over 72 imbalanced datasets with varying types of imbalance.

Beyond Uniform Sampling: Offline Reinforcement Learning with Imbalanced Datasets

Abstract

Offline policy learning is aimed at learning decision-making policies using existing datasets of trajectories without collecting additional data. The primary motivation for using reinforcement learning (RL) instead of supervised learning techniques such as behavior cloning is to find a policy that achieves a higher average return than the trajectories constituting the dataset. However, we empirically find that when a dataset is dominated by suboptimal trajectories, state-of-the-art offline RL algorithms do not substantially improve over the average return of trajectories in the dataset. We argue this is due to an assumption made by current offline RL algorithms of staying close to the trajectories in the dataset. If the dataset primarily consists of sub-optimal trajectories, this assumption forces the policy to mimic the suboptimal actions. We overcome this issue by proposing a sampling strategy that enables the policy to only be constrained to 'good data" rather than all actions in the dataset (i.e., uniform sampling). We present a realization of the sampling strategy and an algorithm that can be used as a plug-and-play module in standard offline RL algorithms. Our evaluation demonstrates significant performance gains in 72 imbalanced datasets, D4RL dataset, and across three different offline RL algorithms. Code is available at https://github.com/Improbable-AI/dw-offline-rl .

1 Introduction

Offline reinforcement learning (RL) [23, 27] aims to learn a decision-making policy that maximizes the expected return (i.e., the sum of rewards over time) using a pre-collected dataset of trajectories, making it appealing for applications where data collection is infeasible or expensive (e.g., recommendation systems [28]). Without loss of generality, it can be assumed that the dataset is generated from an unknown policy π D ( a | s ) , also known as the behavior policy [24]. The goal in offline RL is to learn a policy, π θ ( a | s ) with parameters θ , that exceeds the performance of the behavior policy. In offline RL, a widely recognized issue is the overestimation of Q -values for out-of-distribution state-action pairs, leading to suboptimal policies [8, 20, 22]. This stems from incomplete coverage of the state-action space in the dataset, causing the learning algorithm to consider absent states and actions during optimization.

Most state-of-the-art offline RL algorithms [7, 9, 20, 22, 25] mitigate the issue of OOD Q-values by constraining the distribution of actions of the learned policy π θ ( a | s ) , to be close to the distribution of actions in the dataset. This results in a generic objective with the following form:

where J ( π θ ) denotes the expected return of the policy π θ , D denotes the dataset, C is a regularization term that penalizes the policy π θ for deviating from the state-action pairs in the dataset, and α is the

High-return

action

Low-return

action

Learned

policy

π

θ

Behavior

policy

π

/u1D49F

Re-weighted behavior

policy

π

/u1D49F

w

Action likelihood

Action likelihood

(a) Regularized offline RL with uniform sampling

(b) Regularized offline RL with re-weighted dataset

Action

Action

hyper-parameter balancing the conflicting objectives of maximizing returns while also staying close to the data distribution. This prevents offline RL algorithms from learning behaviors that produce action distributions that diverge significantly from the behavior policy.

An easy-to-understand example of choice for C is the squared distance between the policy and the data [7], C ( s, a ) := ∥ π θ ( s ) -a ∥ 2 2 , where π θ ( s ) denotes the mean of the action distribution π θ ( . | s ) and a is an action sampled from the dataset. When the collected dataset is good, i.e., mostly comprising high-return trajectories, staying close to the data distribution is aligned with the objective of maximizing return, and existing offline RL algorithms work well. However, in scenarios where the dataset is skewed or imbalanced, i.e., contains only a few high-return trajectories and many low-return trajectories, staying close to the data distribution amounts to primarily imitating low-performing actions and is, therefore, detrimental. Offline RL algorithms struggle to learn high-return policies in such scenarios [12]. We present a method that overcomes this fundamental limitation of offline RL. Our method is plug-and-play in the sense that it is agnostic to the choice of the offline RL algorithm.

Our key insight stems from the observation that current methods are unnecessarily conservative by forcing the policy π θ to stay close to all the data. Instead, we would ideally want the policy π θ to be close to the best parts of the offline dataset. This suggests that we should constrain π θ to only be close to state-action pairs that would be generated from a policy that achieves high returns, for instance, the (nearly) optimal policy π ∗ . In offline scenarios, where collecting additional data is prohibited, to mirror the data distribution of π ∗ as much as possible, we can re-weight existing data (i.e., importance sampling [17]). We instantiate this insight in the following way: Represent the distribution induced by a better policy π D w (initially unknown) by re-weighting data points in the dataset with importance weights w ( s, a ) and denote this distribution as D w ( s, a ) . Under this weighting, the offline RL algorithm's training objective can be written as:

Solving for π θ using the re-weighted objective constrains the policy to be close to the better policy π D , and, therefore, allows learning of performant policies. The key challenge is determining the weights w since the state-action distribution of the better policy π D w is initially unknown. To address this, we employ off-policy evaluation techniques [32, 26, 44] to connect the data distribution D w with the expected return of the policy π D w that would generate it. This allows one to optimize the importance weights with respect to its expected return as follows:

Here r ( s, a ) denotes reward of state-action pair ( s, a ) . By exploiting this connection, we optimize the importance weights w to maximize the expected return J ( π D w ) of its corresponding policy π D w subject to necessary constraints (i.e., Bellman flow conservation constraint [35, 32]). This enables us to obtain a better importance weights w .

We evaluate our method with state-of-the-art offline RL algorithms [22, 18] and demonstrate performance gain on 72 imbalanced datasets [12, 6]. Our method significantly outperforms prior work [12]

in challenging datasets with more diverse initial states and fewer trajectories ( 20 × smaller than existing datasets). These datasets pose greater challenges, yet they are crucial for practical applications since real-world datasets are often small and exhibit diverse initial states (e.g., robots with different starting positions, where data comes from human teleoperation).

2 Preliminaries

Typical (online) RL. Reinforcement learning is a formalism that enables us to optimize an agent's policy in a Markov decision process (MDP [35]). The agent (i.e., decision-maker) starts from an initial state s 0 sampled from an initial state distribution ρ 0 ( . ) . At each timestep t , the agent perceives the state s t , takes an action a t ∼ π ( . | s t ) with its policy π , receives a reward r t = r ( s t , a t ) from the environment, and transitions to a next state s t +1 sampled from the environment dynamics T ( . | s t , a t ) until reaching terminal states. The goal in RL is to learn a policy π to maximize the γ -discounted expected, infinite-horizon return J γ ( π ) = E s 0 ∼ ρ 0 ,a t ∼ π ( . | s t ) ,s t +1 ∼T ( . | s t ,a t ) [ ∑ ∞ t =0 γ t r ( s t , a t ) ] . Typical (online) RL algorithms estimate the policy π 's expected return (policy evaluation) from trajectories τ = ( s 0 , a 0 , r 0 , s 1 , a 1 , r 1 · · · ) generated by rolling out π , and update the policy π toward increasing J γ ( π ) (policy improvement), and repeat the processing by performing rollouts with the updated policy.

Offline RL. With no interaction with the environment allowed during the course of learning, offline RL algorithms aim to learn a policy π that maximizes return, entirely using a fixed dataset D that was collected by an arbitrary and unknown 'behavior policy" π D (e.g., humans or pre-programmed controllers). These methods typically aim to estimate the return of a policy π via techniques such as Q-learning or actor-critic, only using batches of state-action pairs ( s t , a t ) uniformly drawn from D . We will denote this estimate value of the return as ̂ J γ D ( π ) . The dataset D consists of N trajectories rolled out by π D :

where T i denotes the length of τ i . In practice, a limit on trajectory length is required since we can only collect finite-length trajectories [34]. When the states that the policy π would encounter and the actions that π would take are not representative in the dataset D , the estimated return ̂ J γ D ( π ) is typically inaccurate [8, 20, 22]. Thus, most offline RL algorithms learn the policy π with pessimistic or conservative regularization that penalizes shift of π from the behavior policy π D that collected the dataset. Typically, implicitly or explicitly, the policy π learned by most offline RL algorithms can be thought of as optimizing the following regularized objective:

where C measures some kind of divergence (e.g., Kullback-Leibler divergence [7]) between π and π D , and α ∈ R + denotes the strength of regularization.

3 Problem Statement: Unnecessary Conservativeness in Imbalanced Datasets

In this section, we describe the issue of offline RL in imbalanced datasets. While algorithms derived from the regularized offline RL objective (Equation 2) attain good performance on several standard benchmarks [6], recent work [12] showed that it leads to 'unnecessary conservativeness" on imbalanced datasets [12] due to the use of constant regularization weight on each state-action pairs ( s, a ) in Equation 2. To illustrate why, we start by defining imbalance of a dataset D using the positive-sided variance of the returns of the dataset (RPSV [12] defined in Definition 3.1). In essence, RPSV measures the dispersion of trajectory returns in the dataset. It indicates the room for improvement of the dataset. Figure 2 illustrates the distribution of trajectory returns in imbalanced datasets with high and low RPSV. Datasets with a low RPSV exhibit a pronounced concentration of returns around the mean value, whereas datasets with a high RPSV display a return distribution that extends away from the mean, towards higher returns. Intuitively, a dataset with high RPSV has trajectories with far higher returns than the average return of the dataset, indicating high chances of finding better data distribution through reweighting. Throughout this paper, we will use the term imbalanced datasets to denote datasets with high RPSV.

Definition 3.1 (Dataset imbalance) . RPSV of a dataset, V + [ G ( τ i )] , corresponds to the secondorder moment of the positive component of the difference between trajectory return: G ( τ i ) := ∑ T i -1 t =0 γ t r ( s i t , a i t ) and its expectation, where τ i denote trajectory in the dataset:

Imbalanced datasets are common in real-world scenarios, as collecting high-return trajectories is often more costly than collecting low-return ones. An example of an imbalanced offline dataset is autonomous driving, where most trajectories are from average drivers, with limited data from very good drivers. Due to the dominance of low-return trajectories, state-action pairs ( s, a ) from these trajectories are oversampled in Equation 2. Consequently, optimizing the regularized objective (Equation 2) would result in a policy that closely imitates the actions from low-performing trajectories that constitute the majority of dataset D , but ideally, we want the policy to imitate actions only on state-action pairs from high-performing trajectories. How-

0.5

0.0

0.5

1.0

1.5

Return

G

(

i

)

0

1

2

3

Density

RPSV

High

Low

ever, current offline RL algorithms [21, 18, 7] use constant regularization weight α (Equation 2). As a result, each state-action pairs is weighted equally, which leads the algorithm to be unnecessarily conservative on all data (i.e., imitating all actions of state-action pairs in the dataset). Further analysis of imbalanced datasets can be found in Appendix A.1.

4 Mitigating Unnecessary Conservativeness By Weighting Samples

In this section, we seek to develop an approach to address the unnecessary conservativeness issue (Section 3) of regularized offline RL algorithms in imbalanced datasets. Adding more experiences from high-performing policies to the dataset would regularize the policy to keep close to highperforming policies and hence easily mitigate the unnecessary conservativeness issue. Though collecting additional experiences (i.e., state-action pairs ( s, a ) ) from the environment is prohibited in offline RL, importance sampling [32] can emulate sampling from another dataset D w since the weighting can be regarded as the density ratio shown below:

where D w ( s, a ) and D ( s, a ) denote the probability density of state-action pairs ( s, a ) in dataset D w and D . Note that D w is unknown but implicitly defined through a given weighting function w . This allows us to adjust the sampling distribution that we train the policy π to, as suggested in the following equivalence:

The remaining question is: how can we determine the weighting w ( s, a ) so that we emulate sampling from a better dataset D w collected by a policy that achieves higher return than the behavior policy π D that collected the original dataset D ?.

4.1 Optimizing the Weightings: Emulating Sampling from High-Performing Policies

Our goal is to discover a weighting function w that can emulate drawing state-action samples from a better dataset D w that is collected by an alternative behavior policy π D w with higher return than the behavior policy π D that collected the original dataset D (i.e., J γ ( π D w ) ≥ J γ ( π D ) ). We make use of density-ratio-based off-policy evaluation methods [32, 29, 44] to determine if a weighting function corresponds to a high-return policy. Note that we do not propose a new off-policy evaluation approach but rather apply the existing off-policy evaluation technique in our problem. By using these

techniques, we can relate the weighting w to the expected return of the alternative behavior policy J ( π D w ) via importance sampling formulation as follows:

In Equation 6, J γ ( π D w ) evaluates the quality of a given weighting function w . It also provides a feasible objective to optimize w , since it only requires obtaining samples from the original dataset D . However, it is important to note that Equation 6 measures γ -discounted return only when the the dataset D w represents a stationary state-action distribution that satisfies Bellman flow conservation constraint [35, 32, 44] in the MDP, as shown in the following equation:

where S and A denote the state and action spaces, respectively, and the discount factor γ determines what discount factor corresponds to in J γ ( π D w ) in Equation 6. We slightly abuse the notation, denoting state marginal as D w ( s ′ ) .

To estimate J γ ( π D w ) from the weighting function w , it is required to impose Bellman flow conservation constraint (Equaton 7) on w . However, it is difficult to impose this constraint due to the dependence of initial state distribution ρ 0 in Equaton 7. Estimating ρ 0 from the first state of each trajectory in the dataset is an option, but it is infeasible when the trajectories do not consistently start from initial states sampled from the distribution ρ 0 . While we could make the assumption that all trajectories begin from initial states sampled from ρ 0 , it would limit the applicability of our method to datasets where trajectories start from arbitrary states. We thus choose not to make this assumption since current offline RL algorithms do not require it.

̸

Instead, since the Bellman flow conservation constraint (Equation 7) only depends on the initial state distribution ρ 0 when γ = 1 , it is possible to bypass this dependence, if we maximize the undiscounted return J ( π D w ) = J γ =1 ( π D w ) (i.e., setting γ = 1 in Equation 6) of the alternative behavior policy π D w . While it deviates from the RL objective presented in Equation 2, undiscounted return is often more aligned with the true objective in various RL applications, as suggested in [13]. Many RL algorithms resort to employing discounted return as an approximation of undiscounted return instead due to the risk of divergence when estimating the undiscounted return using Q-learning [39]. Thus, we constrain the weighting function w to satisfy the Bellman flow conservation constraint with γ = 1 as shown below:

To connect the constraint in Equation 8 to the objective in Equation 6, we rewrite Equation 8 in terms of weightings w according to [29, 32] 2 , as shown below:

where w ( s ) denotes state marginal weighting. Putting the objective (Equation 6) and the constraint (Equation 9) together, we optimize w to maximize the undiscounted expected return of the corresponding alternative behavior policy π D w , as shown in the following:

As the weightings w can be viewed as the density ratio (Equation 4), we call our method as Densityratio Weighting (DW) . We then re-weight offline RL algorithm, as shown in Equation 5. Note that while these weights correspond to ( s, a ) in the dataset, this is sufficient to reweight policy optimization for offline RL.

4.2 Practical Implementation

Optimizing weightings. We begin by addressing the parameterization of the weighting function w ( s, a ) and its state marginal w ( s ′ ) in Equation 10. Though state marginal w ( s ′ ) can derived from

summing w ( s, a ) over action space A , as defined in Equation 9, it can difficult to take summation over a continuous or infinite action space. Thus we opt to parameterize the weightings w ( s, a ) and its state marginal w ( s ′ ) separately. By using the identities [32] D w ( s, a ) = D w ( s ) π D w ( a | s ) and D ( s, a ) = D ( s ) π D ( a | s ) , we can represent w ( s, a ) as the product of two ratios:

Michel et al. [31] showed that ratios can be parameterized by neural networks with exponential output. Thus, we represent state-action weighting w ( s, a ) as w ϕ,ψ ( s, a ) and its state marginal as w ϕ ( s ) , as shown below:

where ϕ and ψ 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 ( D w ||D ) to the objective in Equation 10, as proposed in Zhan et al. [43]. This regularization keeps the state-action distribution D w induced by the learned weighting w close to the original dataset D , preventing w ϕ,ψ ( s, a ) from overfitting to a few rare state-action pairs in 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 ϕ,ψ on violating this constraint in expectation. As a result, we optimize w ϕ,ψ for Equation 10 using stochastic gradient ascent to optimize the following objective (details can be found in Appendix A.3):

where s ′ denotes the next state observed after taking action a at state s , and λ F , λ K ∈ 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 ϕ,ψ 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 ϕ,ψ 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 ̂ J γ D ( π ) and C ( s, a ) at each state-action pairs sampled from the dataset D with w ϕ,ψ ( s, a ) performs better than solely weighting the regularization term C ( s, a ) . For example, when weighting the training objective of implicit Q-learning (IQL) [18] (an offline RL method), the weighted objective J D w ( π ) is: E ( s,a ) ∼D [ w ϕ,ψ ( s, a ) A ( s, a ) log π ( 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 ̂ J γ D ( π ) and regularization C ( s, a ) in the same distribution (i.e., same importance weights) is needed to prevent policy π increasing ̂ J γ D ( π ) by exploiting out-of-distribution actions on states with lower weights w ψ,ϕ ( 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.

5 Experimental Evaluation

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 D

2: Initialize policy π and weighting function w ϕ,ψ

3: while not converged do

4: Sample a batch B of tuples of states, actions, rewards, and next states ( s, a, r, s ′ ) from D

5: Update w ϕ,ψ with batch B using Equation 13

6: Update policy π and value function with an offline RL training objective with weights w ϕ,ψ ( s, a ) and 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 -σ % of trajectories from the random-v2 dataset (low-performing) and σ % 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 -σ % of random and σ % of medium trajectories is denoted as random-mediumσ % . We evaluate our method and the baselines on these imbalanced datasets across four σ ∈ { 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:

where ( s i t , a i t ) denotes the state-action pair at timestep t of trajectory τ i . G K % represents a threshold for selecting the topK % of trajectories, with K chosen from { 10 , 20 , 50 } as practiced in [12]. V ( s i 0 ) denotes the value of the initial state s i 0 in trajectory τ i , and the coefficient η represents the temperature coefficient in a Boltzmann distribution. We consider three levels of η : 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 algorithm 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 AWand 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 η 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

Uniform

PF-10%

PF-20%

PF-50%

AW-L

AW-M

AW-H

AW-XH

DW-AW (ours)

DW-Uniform (ours)

CQL

IQL

TD3BC

0

20

40

60

80

IQM

normalized return

(a) Results on imbalanced datasets of trajectories with similar initial states (Section 5.1).

CQL

IQL

TD3BC

0

10

20

30

40

IQM

normalized return

(b) Results on smaller version of datasets used in Figure 3a.

(AW-M) falls short of matching the performance of our DW-AW, demonstrating the effectiveness of DWin 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 AWhas 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.

6 Related Work

Our approach builds upon recent advances in off-policy evaluation techniques, specifically densityratio importance correction estimation (DiCE) [32]. DiCE has been primarily used for policy evaluation [29, 32, 10], while our method make use 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,

Uniform

PF-10%

PF-20%

PF-50%

AW-L

AW-M

AW-H

DW-AW (ours)

DW-Uniform (ours)

CQL

IQL

TD3BC

0

20

40

60

IQM

normalized return

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 importanceweighted 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 stateof-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.

References

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.

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.

S. P. Boyd and L. Vandenberghe. Convex optimization . Cambridge university press, 2004.

J. Buckman, C. Gelada, and M. G. Bellemare. The importance of pessimism in fixed-dataset policy optimization. arXiv preprint arXiv:2009.06799 , 2020.

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.

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.

S. Fujimoto and S. S. Gu. A minimalist approach to offline reinforcement learning. Advances in neural information processing systems , 34:20132-20145, 2021.

S. Fujimoto, D. Meger, and D. Precup. Off-policy deep reinforcement learning without exploration. arXiv preprint arXiv:1812.02900 , 2018.

S. Fujimoto, H. van Hoof, and D. Meger. Addressing function approximation error in actor-critic methods. arXiv preprint arXiv:1802.09477 , 2018.

S. Fujimoto, D. Meger, and D. Precup. A deep reinforcement learning approach to marginalized importance sampling with the successor representation. In International Conference on Machine Learning , pages 3518-3529. PMLR, 2021.

T. Haarnoja, A. Zhou, K. Hartikainen, G. Tucker, S. Ha, J. Tan, V. Kumar, H. Zhu, A. Gupta, P. Abbeel, et al. Soft actor-critic algorithms and applications. arXiv preprint arXiv:1812.05905 , 2018.

Z.-W. Hong, R. Tachet des Combes, P. Agrawal, and R. Laroche. Harnessing mixed offline reinforcement learning datasets via trajectory reweighting. International Conference on Learning Representations (ICLR) , 2023.

S. Kakade. Optimizing average reward using discounted rewards. In Computational Learning Theory: 14th Annual Conference on Computational Learning Theory, COLT 2001 and 5th European Conference on Computational Learning Theory, EuroCOLT 2001 Amsterdam, The Netherlands, July 16-19, 2001 Proceedings , pages 605-615. Springer, 2001.

D. Kalashnikov, J. Varley, Y. Chebotar, B. Swanson, R. Jonschkowski, C. Finn, S. Levine, and K. Hausman. Mt-opt: Continuous multi-task robotic reinforcement learning at scale. arXiv preprint arXiv:2104.08212 , 2021.

G.-H. Kim, S. Seo, J. Lee, W. Jeon, H. Hwang, H. Yang, and K.-E. Kim. Demodice: Offline imitation learning with supplementary imperfect demonstrations. In International Conference on Learning Representations , 2021.

D. P. Kingma and J. Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.6980 , 2014.

T. Kloek and H. K. Van Dijk. Bayesian estimates of equation system parameters: an application of integration by monte carlo. Econometrica: Journal of the Econometric Society , pages 1-19, 1978.

I. Kostrikov, A. Nair, and S. Levine. Offline reinforcement learning with implicit q-learning. arXiv preprint arXiv:2110.06169 , 2021.

I. Kostrikov, A. Nair, and S. Levine. Offline reinforcement learning with implicit q-learning. In International Conference on Learning Representations , 2022. URL https://openreview. net/forum?id=68n2s9ZJWF8 .

A. Kumar, J. Fu, G. Tucker, and S. Levine. Stabilizing off-policy q-learning via bootstrapping error reduction. URL http://arxiv.org/abs/1906.00949 .

A. Kumar, J. Fu, M. Soh, G. Tucker, and S. Levine. Stabilizing off-policy q-learning via bootstrapping error reduction. Advances in Neural Information Processing Systems , 32, 2019.

A. Kumar, A. Zhou, G. Tucker, and S. Levine. Conservative q-learning for offline reinforcement learning. In Proceedings of the 34th International Conference on Neural Information Processing Systems , NIPS'20, Red Hook, NY, USA, 2020. Curran Associates Inc. ISBN 9781713829546.

S. Lange, T. Gabel, and M. Riedmiller. Batch reinforcement learning. In Reinforcement learning , pages 45-73. Springer, 2012.

R. Laroche and R. T. des Combes. On the occupancy measure of non-Markovian policies in continuous MDPs. In Proceedings of the 40th International Conference on Machine Learning, ICML 2023 , Proceedings of Machine Learning Research. PMLR, 2023.

R. Laroche, P. Trichelair, and R. Tachet des Combes. Safe policy improvement with baseline bootstrapping. In Proceedings of the 36th International Conference on Machine Learning (ICML long oral + EWRL + UAI workshop) , 2019.

J. Lee, W. Jeon, B. Lee, J. Pineau, and K.-E. Kim. Optidice: Offline policy optimization via stationary distribution correction estimation. In International Conference on Machine Learning , pages 6120-6130. PMLR, 2021.

S. Levine, A. Kumar, G. Tucker, and J. Fu. Offline reinforcement learning: Tutorial, review, and perspectives on open problems. arXiv preprint arXiv:2005.01643 , 2020.

L. Li, W. Chu, J. Langford, and X. Wang. Unbiased offline evaluation of contextual-banditbased news article recommendation algorithms. In Proceedings of the fourth ACM international conference on Web search and data mining , pages 297-306, 2011.

Q. Liu, L. Li, Z. Tang, and D. Zhou. Breaking the curse of horizon: Infinite-horizon off-policy estimation. Advances in Neural Information Processing Systems , 31, 2018.

Y. Ma, A. Shen, D. Jayaraman, and O. Bastani. Versatile offline imitation from observations and examples via regularized state-occupancy matching. In K. Chaudhuri, S. Jegelka, L. Song, C. Szepesvari, G. Niu, and S. Sabato, editors, Proceedings of the 39th International Conference on Machine Learning , volume 162 of Proceedings of Machine Learning Research , pages 14639-14663. PMLR, 17-23 Jul 2022. URL https://proceedings.mlr.press/v162/ ma22a.html .

P. Michel, T. Hashimoto, and G. Neubig. Distributionally robust models with parametric likelihood ratios. International Conference on Learning Representations , 2022.

O. Nachum, Y. Chow, B. Dai, and L. Li. Dualdice: Behavior-agnostic estimation of discounted stationary distribution corrections. In Advances in Neural Information Processing Systems , pages 2315-2325, 2019.

O. Nachum, B. Dai, I. Kostrikov, Y. Chow, L. Li, and D. Schuurmans. Algaedice: Policy gradient from arbitrary experience. arXiv preprint arXiv:1912.02074 , 2019.

F. Pardo, A. Tavakoli, V. Levdik, and P. Kormushev. Time limits in reinforcement learning. In International Conference on Machine Learning , pages 4045-4054. PMLR, 2018.

M. L. Puterman. Markov decision processes: discrete stochastic dynamic programming . John Wiley & Sons, 2014.

P. Rashidinejad, B. Zhu, C. Ma, J. Jiao, and S. Russell. Bridging offline reinforcement learning and imitation learning: A tale of pessimism. Advances in Neural Information Processing Systems , 34:11702-11716, 2021.

P. Rashidinejad, H. Zhu, K. Yang, S. Russell, and J. Jiao. Optimal conservative offline rl with general function approximation via augmented lagrangian. arXiv preprint arXiv:2211.00716 , 2022.

A. Singh, A. Kumar, Q. Vuong, Y. Chebotar, and S. Levine. Offline rl with realistic datasets: Heteroskedasticity and support constraints. arXiv preprint arXiv:2211.01052 , 2022.

R. S. Sutton and A. G. Barto. Reinforcement learning: An introduction . 2018.

Q. Vuong, A. Kumar, S. Levine, and Y. Chebotar. Dasco: Dual-generator adversarial support constrained offline reinforcement learning. Advances in Neural Information Processing Systems , 35:38937-38949, 2022.

H. Xu, X. Zhan, H. Yin, and H. Qin. Discriminator-weighted offline imitation learning from suboptimal demonstrations. In International Conference on Machine Learning , pages 2472524742. PMLR, 2022.

T. Yu, A. Kumar, Y. Chebotar, K. Hausman, S. Levine, and C. Finn. Conservative data sharing for multi-task offline reinforcement learning. Advances in Neural Information Processing Systems , 34:11501-11516, 2021.

W. Zhan, B. Huang, A. Huang, N. Jiang, and J. Lee. Offline reinforcement learning with realizability and single-policy concentrability. In Conference on Learning Theory , pages 2730-2775. PMLR, 2022.

R. Zhang, B. Dai, L. Li, and D. Schuurmans. Gendice: Generalized offline estimation of stationary values. arXiv preprint arXiv:2002.09072 , 2020.

A Appendix

A.1 Additional Discussion on Imbalanced Datasets

We consider imbalanced datasets that mixes the trajectories logged by high- and low- performing behavior policies, where trajectories from low-performing behavior policies predominate the dataset. When dealing with such imbalanced datasets, the regularized offline RL objective (Equation 2) tends to constrain the learned policy π to stay close to the overall mediocre behavior policy rather than the high-performing one, which is an unnecessary form of conservativeness. We illustrate this issue from intuitive and analytical aspects in the following.

Intuitively, for example, when considering a dataset D that mix the trajectories from two behavior policies π H D and π L D where π H D significantly outperforms π L D (i.e., J ( π H D ) ≥ J ( π L D ) ), if the dataset is imbalanced such that state-action pairs ( s t , a t ) logged by π L D predominate, the regularized offline RL objective will penalize deviation from π L D more than deviation from π H D . This is because uniform sampling on the imbalanced dataset oversamples state-action pairs ( s t , a t ) from π L D , leading to an over-weighting of the regularization term on those state-action pairs ( s t , a t ) .

Analytically, the issue of unnecessary conservativeness (Section 3) can be explained by how the performance of offline RL J ( π ) is dependent on the performance of the behavior policy J ( π D ) , as suggested in [12]. In the case of an imbalanced dataset where the behavior policy can be regarded a mixture of π H D and π L D , the performance of the behavior policy J ( π D ) can be estimated by the expected return of the distribution of trajectory in the dataset D , as shown below:

where D ( τ i ) denotes the probability mass of trajectory τ i in D , G ( τ i ) := ∑ T i -1 t =0 r ( s i t , a i t ) is the return (i.e., sum of rewards) of trajectory τ i , and T i is its length. Since the trajectories from π L D dominate the dataset, the average return is mostly determined by low-return trajectories, leading to low J ( π D ) . This prevents pessimistic and conservative algorithms [22, 7, 18, 20, 25] from achieving high J ( π ) with low J ( π D ) . Both types of algorithms aim to learn a policy π that outperforms π D (i.e., J ( π ) ≥ J ( π D ) ), while constraining the policy π to stay close to π D [4]. Recent theoretical analysis by Singh et al. [38] on conservative Q-learning (CQL) [22] (an offline RL algorithm) shows that the performance improvements of the learned policy π over the behavior policy is upper-bounded by the deviation of π from π D as shown in: J ( π ) -J ( π D ) ≤ C E ( s,a ) ∼D [ C ( s, a )] , where C is a constant and C denotes the regularization penalty (see Section 2). This implies that regularized objective would constrain the policy π from deviating from π D , thus corroborating the findings in [12] despite the absence general theoretical guarantees for all offline RL algorithms.

Why our method improves performance on imbalanced datasets? Our importance reweighting method can be seen as a modification of the performance of the behavior policy, which represents the lower bound in conservative and pessimistic offline RL algorithms. By adjusting the weights of stateaction pairs sampled from the dataset, we simulate the sampling of data from an alternative dataset D w that would be generated by an alternative behavior policy π D w . By maximizing the expected return J ( π D w ) of this alternative behavior policy π D w , we can obtain a π D w that outperforms the original behavior policy π D , i.e., J ( π D w ) ≥ J ( π D ) . As mentioned earlier, most offline RL algorithms aim to achieve a policy that performs at least as well as the behavior policy that collected the dataset. By training an offline RL algorithm using a dataset D w induced by reweighting with w , we obtain a higher lower bound on the policy's performance, as J ( π D w ) ≥ J ( π D ) .

A.2 Why regularizing the importance weightings with KL divergence does not prevent policy improvement?

In Section 4.2, we employ a KL divergence penalty denoted as D KL ( D||D w ) to impose a penalty on the extent to which the weighted data distribution D w diverges from the original data distribution D . It is important to note that this penalty does not inhibit the learned weightings from potentially improving the data distribution for regularized offline RL algorithms. Essentially, our approach involves optimizing the data distribution for offline RL algorithms first. Subsequently, these offline RL algorithms are subjected to regularization based on this optimized distribution, thereby further

enhancing the policy. This framework enables offline RL algorithms to deviate more significantly from the original data distribution.

To illustrate this concept, consider a scenario where our density-ratio weightings are constrained within an ϵ KL divergence margin with respect to the original data distribution. A regularized offline RL algorithm, which is also subject to an ϵ KL divergence constraint with respect to this weighted data distribution, can potentially improve its performance even when it deviates by more than 2 ϵ from the original data distribution. It is worth noting that some readers might think this as being analogous to reducing the regularization weight α . However, reducing α (Equation 2) lets offline RL algorithms to determine how far they should deviate from the original data distribution on their own. In our experimental investigations, we have observed that reducing α often results in worse performance compared to utilizing our DW method, as evidenced by the results presented in Figure [small kl].

A.3 Implementation details

A.3.1 Offline RL algorithms

We implement our method and other baselines on top of two offline RL algorithms: Conservative Q-Learning (CQL) [22] and Implicit Q-Learning (IQL) [18]. Our baselines, advantage-filtering (AW) and percentage-filtering (PF), are weighted-sampling methods (see Section 5.1). Therefore, they do not require changing the implementation of offline RL algorithms. In the following sections, we explain the details of our importance weighting method for each algorithm.

CQL. We reweight both the actor (policy) and the critic (Q-function) in CQL as follows:

where ¯ H denotes the target entropy used in soft actor critic (SAC) [11] and w ϕ and w ϕ,ψ denote the weights predicted by our method (see Section 4). We follow the notations in CQL paper [22]. The implementation details of computing Equation 17 can be found in [22]. Our implementation is adpated from open-sourced implementation: JaxCQL 3 .

IQL. We reweight state-value function ( V ), state-action value function (i.e,. Q-function, Q ), and policy ( π ) in IQL as follows:

where L τ 2 denotes the upper expectile loss [18] and β denotes the temperature parameter for IQL. Our implementation is adapted from the official implementation 4 for implicit Q-learning (IQL) [18].

We attach our implementation in the supplementary material, where CQL and IQL implementations are in JaxCQL and implicit_q_learning , respectively.

A.3.2 Density-weighting function

The following is the training objective of the importance weighting function w ϕ,ψ and w ϕ in our method, as described in Equation 10 (Section 4.2):

The KL regularization term can be expressed as follows:

Thus, we can rewrite the training objective (Equation 10) to the follows:

The objective in Equation 24 can be optimized by minimizing the loss function L ( ϕ, ψ ) using stochastic gradient descent shown below:

where B denote batch size and ( s i , a i ) denotes the i th state-action pair in the batch. ¯ w ϕ,ψ ( s i , a i ) denotes the batch normalized weights predictions [31] that is defined as follows:

where w ϕ,ψ ( s i , a i ) denotes the importance weights predictions from the network ϕ and ψ (see Section 4.2). As [31] suggests, applying batch-level normalization on the importance weights predictions can make the normalization requirement of importance weights, i.e., E [ w ] = 1 , more likely being satisfied during optimization than other approaches (e.g., adding normalization penalty term) (see [31] for details).

A.4 Evaluation detail

A.4.1 Training details

For both the baselines and our method, we train CQL and IQL for a total of one million gradient steps. The hyperparameters used for CQL and IQL are set to the optimal values recommended in

their publicly available implementations. As for the training of our density weighting functions ϕ, ψ , we employ the same network architecture as the Q-value function architectures in CQL and IQL, which consists of a two-layer Multilayer Perceptron (MLP) with 256 neurons and ReLU activation in each layer. To minimize the objective defined in Equation 25, we train ϕ and ψ using the Adam optimizer [16] with a learning rate of 0.0001 and a batch size of 256. In each gradient step of CQL and IQL, we train ϕ and ψ for one gradient step as well.

We conducted hyperparameter search in the datasets with diverse trajectories. For AW with CQL we searched temperatue η : 0 . 01 (L), 0 . 1 (M, the best from the original paper), 1 . 0 (H), 5 . 0 (XH). For AW with IQL, we searched temperature η : 0 . 01 (L), 0 . 2 (M) (the best from the original paper), 1 . 0 (H), 5 . 0 (XH). For PF in CQL and IQL, we searched K over 0 . 1 , 0 . 2 , and 0 . 5 . The hyperparameter search results are all already presented in Figures 3a. For DW-AW, we use the temperature η used in AW-M for CQL and IQL, except for small datasets. In small datasets (Figure 3b), we use the temperature used in AW-XH for DW-AW. For DW-AW and DW-Uniform, we searched λ K ∈ { 0 . 2 , 1 . 0 } and λ F ∈ { 0 . 1 , 1 . 0 , 5 . 0 } . We use the best found hyperparameter by the time we started the large scale experiments: ( λ K , λ F ) = (0 . 2 , 0 . 1) for CQL and ( λ K , λ F ) = (1 . 0 , 1 . 0) for IQL. We present the hyperparameter search results of DW-Uniform in Figure 2.

A.4.2 Evaluation details

We follow the evaluation protocol used in most offline RL research [22, 7, 18]. Each offline RL algorithm and sampling method (including our DW approach) combination is trained for one million gradient steps with three random seeds in each dataset. We evaluate the policy learned by each method in the environment corresponding to the dataset for 20 episodes every 1000 gradient step. The main performance metric reported is the interquartile mean (IQM) [1] of the normalized mean return over the last 10 rounds of evaluation 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.

Compute. We ran all the experiments using workstations with two RTX 3090 GPUs, AMD Ryzen Threadripper PRO 3995WX 64-Cores CPU, and 256GB RAM.

A.4.3 Dataset curation.

Following the protocol in prior offline RL benchmarking [6, 12], we develop representative datasets for Scenarios (i) and (ii) using the locomotion tasks from the D4RL Gym suite.

Scenario (i): Datasets with trajectories starting from 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. As suggested in [12], these datasets are generated by combining 1 -σ % of trajectories from the random-v2 dataset (low-performing) and σ % 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 -σ % of random and σ % of medium trajectories is denoted as random-mediumσ % . We evaluate our method and the baselines on these imbalanced datasets across four σ ∈ { 1 , 5 , 10 , 50 } , four environments. We construct 32 of this type of datasets in all the combinations of { ant , halfcheetah , hopper , walker2d } × { random-medium , random-expert } × { σ ∈ { 1 , 5 , 10 , 50 } }. Also, 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. As every method fails to suprpass random policy when σ = 1 and σ = 5 , we only evaluate all the methods in σ = 10 . Note that the results in σ = 50 is not different from the results in larger version of mixed dataset with σ = 50 since the amount of high-performing trajectories is still sufficient when σ = 50 .

Scenario (ii): Datasets with trajectories starting from diverse initial states. Trajectories in this type of dataset start from a wider range of initial states and have varying lengths. This characteristic is designed to simulate scenarios where trajectories from different parts of a task are available in the dataset. 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, such as merging onto a highway, changing lanes, navigating through intersections, and parking, 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. We curate this type of datasets by adapting the datasets from Scenario (i). These datasets are created by combining 1 -σ % of trajectory segments from the random-v2 dataset (representing low-performing policies) and σ % of trajectory segments from either the medium-v2 or expert-v2 dataset (representing high-performing policies) for each locomotion environment in the D4RL benchmark. Each trajectory segment is a subsequence of a trajectory in the dataset of Scenario (i). Specifically, for each trajectory τ i , a trajectory segment is selected, ranging in length from 10 to 50 timesteps within τ i . We chose the minimum and maximum lengths to be 10 and 50 timesteps, respectively, based on our observation that most locomotion behaviors exhibit intervals lasting between 10 and 50 timesteps. Since locomotion behaviors are periodic in nature, these datasets can simulate locomotion recordings from various initial conditions, such as different poses and velocities.

The datasets in Scenario (ii) capture a different form of imbalance that is overlooked in the datasets of Scenario (i). The imbalanced datasets in Scenario (i) combine full trajectories from low- and high-performing policies, starting from similar initial states. These datasets capture the imbalance resulting from the policies themselves, while overlooking the imbalance caused by different initial conditions. Even when using the same behavior policy, the initial conditions, i.e., the initial states of trajectories, can lead to an imbalance in the dataset, as certain initial conditions tend to yield higher returns compared to others. For example, an agent starting near the goal and another starting far away from the goal may achieve significantly different returns, even when following the same policy. This type of imbalance can exist in real-world datasets if certain initial conditions are oversampled during the dataset collection process.

Figure 5 presents the distribution of normalized returns for both types of datasets in both scenarios. The trajectory returns are normalized to a range of 0 to 1 using max-min normalization, specifically ( G ( τ i ) -G min ) / ( G max -G min ) , where G min and G max denote the minimum and maximum trajectory returns in the dataset, respectively. From Figure 5, we observe that the two types of datasets exhibit different return distributions. The return distribution in Scenario (i) (i.e., mixed) is closer to a bimodal distribution, where a trajectory is either from a high- or low-performing policy. On the other hand, the return distribution in Scenario (ii) (i.e., mixed (diverse)) follows a heavy-tailed distribution, where high-return trajectories are located in the long tails. This indicates that in addition to the imbalance resulting from the combination of trajectories from low- and high-performing policies, the presence of diverse initial states introduces another type of imbalance in the dataset, where initial states that easily lead to high returns are much less prevalent than others.

A.5 Additional results

We present the full results in total of 113 datasets in Tables 4 and 5, including original D4RL datasets [6], mixed datasets (Figure 3a, denoted as Mixed), mixed datasets with diverse initial states (Figure 4, denoted as Mixed (diverse)), and small datasets (Figure 3b, denoted as Mixed (small)). As [12] suggested, the original D4RL datasets are not imbalanced and thus weighted sampling and importance weighting methods perform on par with uniform sampling.

A.5.1 Four room example for trajectory stitching

To assess DW's trajectory stitching ability, we conducted an experiment in a didactic four-room environment [26].

Experiment setup: See Figure 6 for the environment illustration. The agent starts from an orange initial state, traverses non-red cells, gaining +1 reward at the green goal, else zero. To test trajectory stitching, a suboptimal dataset with 1000 trajectories was generated, where none of each is optimal trajectory. Due to absence of optimal trajectories, up-weighting trajectories (i.e., AW and PF) with high-returns won't produce a state-action distribution matching the optimal policy. Thus, if a method can generate state-action distribution matching the one of the optimal policy, it indicates that the method is able to stitch trajectories because it can identify optimal state-action pairs leading to the goal even though those optimal state-action pairs are not observed in the same trajectory in the dataset.

0.00

0.25

0.50

0.75

1.00

0

10

20

30

Density

ant-random-medium-0.01-v2

mixed

mixed (diverse)

0.00

0.25

0.50

0.75

1.00

0

10

20

ant-random-medium-0.05-v2

mixed

mixed (diverse)

0.00

0.25

0.50

0.75

1.00

0

5

10

15

ant-random-medium-0.1-v2

mixed

mixed (diverse)

0.0

0.5

1.0

0.0

2.5

5.0

7.5

ant-random-medium-0.5-v2

mixed

mixed (diverse)

0.00

0.25

0.50

0.75

1.00

0

20

40

Density

ant-random-expert-0.01-v2

mixed

mixed (diverse)

0.00

0.25

0.50

0.75

1.00

0

10

20

ant-random-expert-0.05-v2

mixed

mixed (diverse)

0.00

0.25

0.50

0.75

1.00

0

10

20

ant-random-expert-0.1-v2

mixed

mixed (diverse)

0.0

0.5

1.0

0

5

10

ant-random-expert-0.5-v2

mixed

mixed (diverse)

0.00

0.25

0.50

0.75

1.00

0

25

50

75

Density

hopper-random-medium-0.01-v2

mixed

mixed (diverse)

0.00

0.25

0.50

0.75

1.00

0

20

40

60

hopper-random-medium-0.05-v2

mixed

mixed (diverse)

0.00

0.25

0.50

0.75

1.00

0

20

40

60

hopper-random-medium-0.1-v2

mixed

mixed (diverse)

0.00

0.25

0.50

0.75

1.00

0

10

20

hopper-random-medium-0.5-v2

mixed

mixed (diverse)

0.00

0.25

0.50

0.75

1.00

0

50

100

Density

hopper-random-expert-0.01-v2

mixed

mixed (diverse)

0.00

0.25

0.50

0.75

1.00

0

25

50

75

hopper-random-expert-0.05-v2

mixed

mixed (diverse)

0.00

0.25

0.50

0.75

1.00

0

20

40

60

hopper-random-expert-0.1-v2

mixed

mixed (diverse)

0.00

0.25

0.50

0.75

1.00

0

10

20

30

hopper-random-expert-0.5-v2

mixed

mixed (diverse)

0.0

0.5

1.0

0

10

20

30

Density

halfcheetah-random-medium-0.01-v2

mixed

mixed (diverse)

0.0

0.5

1.0

0

10

20

halfcheetah-random-medium-0.05-v2

mixed

mixed (diverse)

0.0

0.5

1.0

0

5

10

15

halfcheetah-random-medium-0.1-v2

mixed

mixed (diverse)

0.0

0.5

1.0

0

2

4

6

halfcheetah-random-medium-0.5-v2

mixed

mixed (diverse)

0.0

0.5

1.0

0

20

40

Density

halfcheetah-random-expert-0.01-v2

mixed

mixed (diverse)

0.0

0.5

1.0

0

10

20

30

halfcheetah-random-expert-0.05-v2

mixed

mixed (diverse)

0.0

0.5

1.0

0

10

20

halfcheetah-random-expert-0.1-v2

mixed

mixed (diverse)

0.0

0.5

1.0

0

2

4

6

halfcheetah-random-expert-0.5-v2

mixed

mixed (diverse)

0.00

0.25

0.50

0.75

1.00

0

50

100

150

Density

walker2d-random-medium-0.01-v2

mixed

mixed (diverse)

0.00

0.25

0.50

0.75

1.00

0

50

100

walker2d-random-medium-0.05-v2

mixed

mixed (diverse)

0.00

0.25

0.50

0.75

1.00

0

25

50

75

walker2d-random-medium-0.1-v2

mixed

mixed (diverse)

0.00

0.25

0.50

0.75

1.00

0

10

20

walker2d-random-medium-0.5-v2

mixed

mixed (diverse)

0.00

0.25

0.50

0.75

1.00

Normalized return

0

50

100

150

Density

walker2d-random-expert-0.01-v2

mixed

mixed (diverse)

0.00

0.25

0.50

0.75

1.00

Normalized return

0

50

100

walker2d-random-expert-0.05-v2

mixed

mixed (diverse)

0.00

0.25

0.50

0.75

1.00

Normalized return

0

20

40

60

walker2d-random-expert-0.1-v2

mixed

mixed (diverse)

0.00

0.25

0.50

0.75

1.00

Normalized return

0

10

20

walker2d-random-expert-0.5-v2

mixed

mixed (diverse)

Results: Figures 6a, 6b, and 6e display the state-action distributions of behavior policy, optimal policy, and DW; the number above each plot is expected return under the state-action distribution. We see that both AW (Figure 6c) and PF (Figure 6d fail to match the state-action distribution of the optimal policy in Figure 8(b), hence leading to suboptimal performance. In contrast, DW successfully approximates the optimal policy's state-action distribution, confirming DW can identify optimal state-action pairs observed in different suboptimal trajectories and stitch them to optimal trajectories.

A.5.2 Comparison with OptDICE

OptDiCE [26] and AlgaeDiCE [33], as well as our method, all involve learning importance weights for policy optimization. However, the usage and learning of importance weights vary across these

J

=0.17

(a) Behavior policy

J

=0.75

(b) Optimal policy

J

=0.35

(c) PF

J

=0.37

(d) AW

J

=0.74

(e) DW (ours)

methods. We compare our method with OptDiCE as it is the state-of-the-art method on policy optimization using DiCE. In the following, we illustrate the difference between our method and OptDiCE.

Learning importance weights. OptDiCE learns the importance weights through solving the optimization problem same as Equation 10 in an approach different from ours. OptDiCE learns the importance weights through optimizing the following primal-dual objective:

where f can be any convex function (e.g., f ( x ) := x log x ), ˆ e ν ( s, a, s ′ ) := r ( s, a ) + γν ( s ′ ) -ν ( s ) , and x + := max(0 , x ) . The coefficient α denotes the strength of regularization: the higher the α is, the uniform the importance weights are. The learned importance weights w ( s, a ) are expressed in terms of ν as follows:

On the other hand, our method learns the importance weights without solving the min-max optimization.

Applying importance weights. OptDiCE extracts the policy from the learned importance weights using information projection (I-projection) [26]. The implementation of I-projection is non-trivial, while its idea is close to weighted behavior cloning objective shown as follows:

Differing from OptDiCE, we apply the importance weights to actor-critic offline RL algoritmhs like CQL [22] and IQL [18].

We present the performance of our methods in CQL and IQL, and OptDiCE in imbalanced datasets used in Scenario (i) in Figure 7, showing that OptDiCE underpeform all the other methods on imbalanced datasets and even underperforms uniform sampling approach. This result indicates that even though OptDiCE learns the importance weights by solving a similar objective with our method, OptDiCE is not effective on imbalanced datasets.

A.5.3 Comparison with weights learned by OptDiCE

While the policy learned using OptDiCE may not be effective on imbalanced datasets, it remains uncertain whether the importance weights learned with OptDiCE can enhance the performance of

CQL

IQL

0

20

40

60

80

IQM

normalized return

Uniform

PF-10%

AW

OptDiCE

DW-AW (ours)

DW-Uniform (ours)

(*: training gets terminated due to NaN values)

offline RL algorithms. To investigate this hypothesis, we reweight the training objective of CQL (Equation 17) using the importance weights learned with OptDiCE (Equation 32). We present a comparison with our method in Table 1, where OptDiCEW refers to CQL trained with OptDiCE weights, and α represents the regularization strength (Equation 30). We examine different values of α ranging from 0 . 1 to 5 . 0 , but none of them consistently outperforms our method. It is worth noting that each configuration results in a significant performance loss in certain datasets. While this observation suggests that the importance weights learned using OptDiCE are not effective in improving the performance of offline RL algorithms, it is important to note that our work does not aim to propose a new off-policy evaluation method based on DiCE. Therefore, a comprehensive comparison between different methods for learning importance weights is left for future research endeavors.

A.5.4 Hyperparameter studies

In this section, we investigate the hyperparameters of KL regularization strength ( λ K ) and Bellman flow conservation strength ( λ F ). We present the average returns obtained in four specific imbalanced datasets in Table 2. These datasets are chosen because they exhibit significant performance differences between uniform sampling and other methods. Our observations indicate that higher penalties for flow conservation and lower strengths of KL regularization tend to improve performance in both CQL and IQL.

A.5.5 Comparison of re-weighting both objectives and only regularization objective

In theory, the regularization term E ( s,a ) ∼D [ C ( s, a )] (Equation 2) is the only component that can potentially harm performance on imbalanced datasets, as it encourages the policy to imitate the suboptimal actions that dominate the dataset. However, our findings suggest that reweighting all the training objectives in offline RL algorithms leads to improved performance. In this section, we

(b) IQL

compare the performance of CQL trained with reweighted regularization only (referred to as "Reg. only") and reweighting all objectives (referred to as "All"). For the "All" approach, we train CQL using Equation 17. On the other hand, for the "Reg. only" approach, we remove the term w ϕ,ψ ( s, a ) from the α training and the term w ϕ ( s ) from the π training in Equation 17. Table 3 presents the mean return of both approaches, indicating that "All" exhibits slightly better performance compared to "Reg. only".

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n1. Overall good writing quality and clarity.\n2. The proposed approach is well-motivated by adopting existing approaches. \n3. Sufficient technical details for reproducing the experiments.\n\nWEAKNESSES:\n1. 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. \n2. Lack of details for the dataset setup. \n3. 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. \n3. 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.\n\nQUESTIONS:\n4. 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.\n5. 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?\n7. 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? \n8. 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?\n\n[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.\n\n[2] Brandfonbrener D, Whitney W F, Ranganath R, et al. Quantile filtered imitation learning[J]. arXiv preprint arXiv:2112.00950, 2021.\n\n[3] Wu Y, Zhai S, Srivastava N, et al. Uncertainty weighted actor-critic for offline reinforcement learning[J]. arXiv preprint arXiv:2105.08140, 2021.\n\n[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.\n\n[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.\n\nLIMITATIONS:\nSee questions and weaknesses.\n\n---\nraise the score to 6 after reviewing the response from the authors.\n\n---\n\n**Minors**:\n1. 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].\n2. IQL training object for policy improvement (line 234) is wrong. \n3. 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.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-06-10"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nAuthors 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.\n\nSTRENGTHS:\nMethod 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.\n\nWEAKNESSES:\nWhile 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.\n\nQUESTIONS:\nWhat 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.\n\nHow does the method affect algorithms time needed for training?\n\nLIMITATIONS:\nBeside 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.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-06-18"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n- 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.\n- It demonstrates overall performance gains over SoTA offline RL algorithms and other baselines on imbalanced datasets with varying types and degrees of imbalance.\n- Its written style is clear and easy to follow.\n\nWEAKNESSES:\n- The theoretical analysis is relatively weak.\n- Slightly lack of discussion with a type of offline RL algorithms which filter out low-performance trajectories.\n\nQUESTIONS:\n- 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.\n- 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.\n\n[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.\n\n[2] Liu, Minghuan, et al. \"Curriculum offline imitating learning.\" Advances in Neural Information Processing Systems 34 (2021): 6266-6277.\n\nLIMITATIONS:\nYes.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-06"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nConsidering 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).\n\nThis 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\n\nSTRENGTHS:\nI 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.\n\nWEAKNESSES:\nThere 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.\n\nI 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).\n\nQUESTIONS:\nCould 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?\n\nLIMITATIONS:\nSeems fine.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-06"}