SlowGuess/NeurIPS_2023_OCR huggingface dataset:
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(\tau_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.
Figure 5: Return distributions of mixed datasets with similar initial states (blue, Scenario (i)) and diverse initial states (orange, Scenario (ii)). Both types of datasets lead to different kinds of imbalance and return distributions. See Section A.4.3 for details.
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
(a) Behavior policy
(b) Optimal policy
(c) PF
Figure 6: The stationary state-action distributions of behavior policy (a), the optimal policy (b), percentage filtering (i.e., top  $\mathrm{K\%}$ ) (c), advantage weighting (d), and density weighting (e).  $J$  denotes the expected return for each. The agent starts from the orange state and will receive the rewards at the green state (goal) and episode termination. Figure 6a is the empirical state-action distribution formed by 1000 suboptimal trajectories without optimal trajectory to the goal. Reweighting trajectories (i.e., AW and PF) cannot produce the optimal state-action distribution (Figure 6b) in the absence of optimal trajectories in the dataset. In contrast, since DW assign weights to transitions rather than trajectories, it can up-weight good transitions in suboptimal trajectories and stitch them to optimal state-action distribution. As we see, only DW (Figure 6e) successfully matches the optimal policy. This indicates that DW has the ability to stitch trajectories from suboptimal data distribution like Figure 6a while the other methods cannot.
(d) AW
(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:
$$
\min  _ {\nu} \max  _ {w \geq 0} \mathbb {E} _ {(s, a) \sim \mathcal {D}} \left[ e _ {\nu} (s, a) w (s, a) - \alpha f (w (s, a)) \right] + (1 - \gamma) \mathbb {E} _ {s _ {0} \sim p _ {0}} [ \nu (s _ {0}) ] \approx \tag {30}
$$
$$
\min  _ {\nu} \mathbb {E} _ {(s, a, s ^ {\prime}) \sim \mathcal {D}} \left[ \hat {e} _ {\nu} (s, a, s ^ {\prime}) \left(f ^ {\prime}\right) ^ {- 1} \left(\frac {1}{\alpha} \hat {e} _ {\nu} (s, a, s ^ {\prime})\right) _ {+} - \alpha f \left(\left(f ^ {\prime}\right) ^ {- 1} \left(\frac {1}{\alpha} \hat {e} _ {\nu} (s, a, s ^ {\prime})\right) _ {+}\right) \right] + \tag {31}
$$
$$
(1 - \gamma) \mathbb {E} _ {s _ {0} \sim p _ {0}} [ \nu (s _ {0}) ],
$$
where  $f$  can be any convex function (e.g.,  $f(x) \coloneqq x \log x$ ),  $\hat{e}_{\nu}(s, a, s') \coloneqq r(s, a) + \gamma \nu(s') - \nu(s)$ , and  $x_{+} \coloneqq \max(0, x)$ . The coefficient  $\alpha$  denotes the strength of regularization: the higher the  $\alpha$  is, the uniform the importance weights are. The learned importance weights  $w(s, a)$  are expressed in terms of  $\nu$  as follows:
$$
w (s, a) = \max  \left(0, (f) ^ {- 1} \left(\frac {\hat {e} _ {\nu} (s , a , s ^ {\prime})}{\alpha}\right)\right). \tag {32}
$$
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:
$$
\max  _ {\pi} \mathbb {E} _ {(s, a) \sim \mathcal {D} _ {w}} [ \log \pi (a | s) ] = \max  _ {\pi} \mathbb {E} _ {(s, a) \sim \mathcal {D}} [ w (s, a) \log \pi (a | s) ]. \tag {33}
$$
Differing from OptDiCE, we apply the importance weights to actor-critic offline RL algorithms 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 underperforms 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
Figure 7: Results on imbalanced datasets of trajectories with similar initial states (Scenario (i) in Section 5.1). OptDiCE [26] fails to achieve high return on imbalanced datasets and even performs worse than CQL and IQL with uniform sampling. Note that OptDiCE is a density-ratio importance correction estimation (DiCE) based offline RL algorithm rather than a sampling or importance weighting method. While both bars of OptDiCE in this figure denote the same result, we plot the performance of OptDiCE alongside CQL and IQL with different sampling (or weighting) methods for comparison.
OptDiCEW (α = 0.1)	OptDiCEW (α = 1.0)	OptDiCEW (α = 5.0)
hopper-random-expert-diverse-5%-v2	3.3 (-17.5)	39.1 (+18.3)	4.0 (-16.8)
hopper-random-expert-diverse-10%-v2	1.5 (-56.4)	29.7 (-28.2)	11.9 (-46.0)
hopper-random-medium-diverse-5%-v2	-*	6.4 (-58.5)	57.3 (-7.6)
hopper-random-medium-diverse-10%-v2	2.2 (-52.9)	30.1 (-25.0)	41.0 (-14.1)
(*: training gets terminated due to NaN values)
Table 1: Comparison of the weights learned by our method and that learned by OptDiCE approach. Each cell in the table denotes the mean return obtained by CQL trained with OptDiCE weights and the number in the parenthesis indicates the improvements (i.e.,  $\text{score}_{\text{OptDiCE}} - \text{score}_{\text{Ours}}$ ) over the performance of DW-Uniform reported in Figure 4. It shows that OptDiCE underperforms our method in most datasets (i.e., negative improvements), and is sensitive to the regularization strength parameter  $\alpha$  (see Section A.5.3).
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  $\alpha$  represents the regularization strength (Equation 30). We examine different values of  $\alpha$  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  $(\lambda_{K})$  and Bellman flow conservation strength  $(\lambda_{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  $\mathbb{E}_{(s,a)\sim \mathcal{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
(λK, λF)	(0.2, 0.1)	(0.2, 1.0)	(0.2, 5.0)	(1.0, 0.1)	(1.0, 1.0)	(1.0, 5.0)
hopper-random-expert-diverse-5%-v2	20.4	13.8	22.1	9.6	7.5	7.9
hopper-random-expert-diverse-10%-v2	51.7	27.8	42.1	10.9	9.8	5.3
hopper-random-medium-diverse-5%-v2	64.7	66.1	60.3	16.2	5.5	25.4
hopper-random-medium-diverse-10%-v2	55.4	58.9	65.3	3.6	7.2	7.3
(a) CQL
(λK, λF)	(0.2, 0.1)	(0.2, 1.0)	(0.2, 5.0)	(1.0, 0.1)	(1.0, 1.0)	(1.0, 5.0)
hopper-random-expert-diverse-5%-v2	14.7	56.4	54.2	16.8	68.4	75.1
hopper-random-expert-diverse-10%-v2	16.2	67.8	51.0	78.5	63.0	67.5
hopper-random-medium-diverse-5%-v2	53.8	50.8	51.5	50.9	49.5	52.6
hopper-random-medium-diverse-10%-v2	43.4	48.8	53.0	42.5	53.2	50.8
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_{\phi, \psi}(s, a)$  from the  $\alpha$  training and the term  $w_{\phi}(s)$  from the  $\pi$  training in Equation 17. Table 3 presents the mean return of both approaches, indicating that "All" exhibits slightly better performance compared to "Reg. only".
(b) IQL
Table 2: Hyperparameter studies of our DW method in (a) CQL and (b) IQL.  $\lambda_{K}$  and  $\lambda_{F}$  denote the strength of KL-regularization and Bellman flow conservation penalty, respectively (see Section A.3). Lower KL-regularization strength  $\lambda_{F}$  and higher flow conservation penalty strength  $\lambda_{F}$  lead to better performance.
Reg. only	All
hopper-random-expert-diverse-10%-v2	57.1	64.9
walker2d-random-expert-diverse-10%-v2	1.4	8.2
Table 3: Reweighting all terms in Equation 17 shows better performance than reweighting only the regularization term. See Section A.5.5.
Table 4: Full results of average returns of CQL in total of 113 datasets.
Table 5: Full results of average returns of IQL in total of 113 datasets.
Table 6: Full results of average returns of TD3BC in total of 113 datasets.