SlowGuess/NeurIPS_2023_OCR huggingface dataset:
[29] 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.
[30] 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.
[31] P. Michel, T. Hashimoto, and G. Neubig. Distributionally robust models with parametric likelihood ratios. International Conference on Learning Representations, 2022.
[32] 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.
[33] 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.
[34] 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.
[35] M. L. Puterman. Markov decision processes: discrete stochastic dynamic programming. John Wiley & Sons, 2014.
[36] 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.
[37] 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.
[38] 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.
[39] R. S. Sutton and A. G. Barto. Reinforcement learning: An introduction. 2018.
[40] 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.
[41] H. Xu, X. Zhan, H. Yin, and H. Qin. Discriminator-weighted offline imitation learning from suboptimal demonstrations. In International Conference on Machine Learning, pages 24725-24742. PMLR, 2022.
[42] 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.
[43] 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.
[44] R. Zhang, B. Dai, L. Li, and D. Schuurmans. Gendice: Generalized offline estimation of stationary values. arXiv preprint arXiv:2002.09072, 2020.
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  $\pi$  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  $\mathcal{D}$  that mix the trajectories from two behavior policies  $\pi_{\mathcal{D}}^{H}$  and  $\pi_{\mathcal{D}}^{L}$  where  $\pi_{\mathcal{D}}^{H}$  significantly outperforms  $\pi_{\mathcal{D}}^{L}$  (i.e.,  $J(\pi_{\mathcal{D}}^{H}) \geq J(\pi_{\mathcal{D}}^{L}))$ , if the dataset is imbalanced such that state-action pairs  $(s_t, a_t)$  logged by  $\pi_{\mathcal{D}}^{L}$  predominate, the regularized offline RL objective will penalize deviation from  $\pi_{\mathcal{D}}^{L}$  more than deviation from  $\pi_{\mathcal{D}}^{H}$ . This is because uniform sampling on the imbalanced dataset oversamples state-action pairs  $(s_t, a_t)$  from  $\pi_{\mathcal{D}}^{L}$ , 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(\pi)$  is dependent on the performance of the behavior policy  $J(\pi_{\mathcal{D}})$ , as suggested in [12]. In the case of an imbalanced dataset where the behavior policy can be regarded a mixture of  $\pi_{\mathcal{D}}^{H}$  and  $\pi_{\mathcal{D}}^{L}$ , the performance of the behavior policy  $J(\pi_{\mathcal{D}})$  can be estimated by the expected return of the distribution of trajectory in the dataset  $\mathcal{D}$ , as shown below:
$$
J \left(\pi_ {\mathcal {D}}\right) \approx \mathbb {E} _ {\tau_ {i} \sim \mathcal {D}} \left[ G \left(\tau_ {i}\right) \right] = \sum_ {i = 1} ^ {N} \mathcal {D} \left(\tau_ {i}\right) G \left(\tau_ {i}\right), \tag {16}
$$
where  $\mathcal{D}(\tau_i)$  denotes the probability mass of trajectory  $\tau_{i}$  in  $\mathcal{D}$ ,  $G(\tau_i)\coloneqq \sum_{t = 0}^{T_i - 1}r(s_t^i,a_t^i)$  is the return (i.e., sum of rewards) of trajectory  $\tau_{i}$ , and  $T_{i}$  is its length. Since the trajectories from  $\pi_{\mathcal{D}}^{L}$  dominate the dataset, the average return is mostly determined by low-return trajectories, leading to low  $J(\pi_{\mathcal{D}})$ . This prevents pessimistic and conservative algorithms [22, 7, 18, 20, 25] from achieving high  $J(\pi)$  with low  $J(\pi_{\mathcal{D}})$ . Both types of algorithms aim to learn a policy  $\pi$  that outperforms  $\pi_{\mathcal{D}}$  (i.e.,  $J(\pi)\geq J(\pi_{\mathcal{D}})$ ), while constraining the policy  $\pi$  to stay close to  $\pi_{\mathcal{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  $\pi$  over the behavior policy is upper-bounded by the deviation of  $\pi$  from  $\pi_{\mathcal{D}}$  as shown in:  $J(\pi) - J(\pi_{\mathcal{D}})\leq C\mathbb{E}_{(s,a)\sim \mathcal{D}}[\mathcal{C}(s,a)]$ , where  $C$  is a constant and  $\mathcal{C}$  denotes the regularization penalty (see Section 2). This implies that regularized objective would constrain the policy  $\pi$  from deviating from  $\pi_{\mathcal{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 state-action pairs sampled from the dataset, we simulate the sampling of data from an alternative dataset  $\mathcal{D}_w$  that would be generated by an alternative behavior policy  $\pi_{\mathcal{D}_w}$ . By maximizing the expected return  $J(\pi_{\mathcal{D}_w})$  of this alternative behavior policy  $\pi_{\mathcal{D}_w}$ , we can obtain a  $\pi_{\mathcal{D}_w}$  that outperforms the original behavior policy  $\pi_{\mathcal{D}}$ , i.e.,  $J(\pi_{\mathcal{D}_w}) \geq J(\pi_{\mathcal{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  $\mathcal{D}_w$  induced by reweighting with  $w$ , we obtain a higher lower bound on the policy's performance, as  $J(\pi_{\mathcal{D}_w}) \geq J(\pi_{\mathcal{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}(\mathcal{D}||\mathcal{D}_w)$  to impose a penalty on the extent to which the weighted data distribution  $\mathcal{D}_w$  diverges from the original data distribution  $\mathcal{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  $\epsilon$  KL divergence margin with respect to the original data distribution. A regularized offline RL algorithm, which is also subject to an  $\epsilon$  KL divergence constraint with respect to this weighted data distribution, can potentially improve its performance even when it deviates by more than  $2\epsilon$  from the original data distribution. It is worth noting that some readers might think this as being analogous to reducing the regularization weight  $\alpha$ . However, reducing  $\alpha$  (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  $\alpha$  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:
$$
\max  _ {\pi} \mathbb {E} _ {s \sim \mathcal {D}, a \sim \pi (.) | s)} \left[ w _ {\phi} (s) \left(Q (s, a) - \log \pi (a | s)\right) \right] \tag {17}
$$
$$
\min  _ {Q} \alpha \mathbb {E} _ {(s, a) \sim \mathcal {D}} \left[ w _ {\phi , \psi} (s, a) \left(\log \sum_ {a ^ {\prime} \in \mathcal {A}} \exp Q (s, a ^ {\prime}) - Q (s, a)\right) \right] + \tag {18}
$$
$$
\left. \mathbb {E} _ {(s, a, s ^ {\prime}) \sim \mathcal {D}, a ^ {\prime} \sim \pi (.) | s ^ {\prime}} \left[ w _ {\phi , \psi} (s, a) (r (s, a) + \gamma Q (s ^ {\prime}, a ^ {\prime}) - Q (s, a)) ^ {2} \right] \right.
$$
$$
\min  _ {\alpha} \mathbb {E} _ {s \sim \mathcal {D}, a \sim \pi (.) | s} \left[ w _ {\phi} (s) (- \log \pi (a | s) - \bar {\mathcal {H}}) \right] \quad (\text {E n t r o p y c o e f f i c i e n t [ 1 1 ]}), \tag {19}
$$
where  $\bar{\mathcal{H}}$  denotes the target entropy used in soft actor critic (SAC) [11] and  $w_{\phi}$  and  $w_{\phi, \psi}$  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 adapted from open-sourced implementation: JaxCQL<sup>3</sup>.
IQL. We reweight state-value function  $(V)$ , state-action value function (i.e., Q-function,  $Q$ ), and policy  $(\pi)$  in IQL as follows:
$$
\min  _ {V} \mathbb {E} _ {(s, a) \sim \mathcal {D}} \left[ w _ {\phi} (s) L _ {2} ^ {\tau} \left(Q (s, a) - V (s)\right) \right] \tag {20}
$$
$$
\min  _ {Q} \mathbb {E} _ {(s, a, s ^ {\prime}) \sim \mathcal {D}} \left[ w _ {\phi , \psi} (s, a) \left(r (s, a) + \gamma V \left(s ^ {\prime}\right) - V (s)\right) ^ {2} \right] \tag {21}
$$
$$
\max  _ {\pi} \mathbb {E} _ {(s, a) \sim \mathcal {D}} \left[ w _ {\phi , \psi} (s, a) \exp \left(\beta \left(Q (s, a) - V (s)\right)\right) \log \pi (a | s) \right], \tag {22}
$$
where  $L_2^{\tau}$  denotes the upper expectile loss [18] and  $\beta$  denotes the temperature parameter for IQL. Our implementation is adapted from the official implementation<sup>4</sup> 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.
<sup>3</sup>https://github.com/young-geng/JaxCQL
<sup>4</sup>https://github.com/ikostrikov/implicit_q_learning
A.3.2 Density-weighting function
The following is the training objective of the importance weighting function  $w_{\phi, \psi}$  and  $w_{\phi}$  in our method, as described in Equation 10 (Section 4.2):
$$
\max  _ {\phi , \psi} \mathbb {E} _ {(s, a, s ^ {\prime}) \sim \mathcal {D}} \left[ \underbrace {w _ {\phi , \psi} (s , a) r (s , a)} _ {\text {R e t u r n}} - \lambda_ {F} \underbrace {\left(w _ {\phi} \left(s ^ {\prime}\right) - w _ {\phi , \psi} (s , a)\right) ^ {2}} _ {\text {B e l l m a n f l o w c o n s e r v a t i o n p e n a l t y}} \right] - \lambda_ {K} \underbrace {D _ {K L} (\mathcal {D} _ {w} | | \mathcal {D})} _ {\text {K L r e g u l a r i z a t i o n}}.
$$
The KL regularization term can be expressed as follows:
$$
\begin{array}{l} D _ {K L} \left(\mathcal {D} _ {w} \| \mathcal {D}\right) = \sum_ {s, a} \mathcal {D} _ {w} (s, a) \log \frac {\mathcal {D} _ {w} (s , a)}{\mathcal {D} (s , a)} \tag {23} \\ = \sum_ {s, a} \mathcal {D} (s, a) \frac {\mathcal {D} _ {w} (s , a)}{\mathcal {D} (s , a)} \log \frac {\mathcal {D} _ {w} (s , a)}{\mathcal {D} (s , a)} \\ = \mathbb {E} _ {(s, a) \sim \mathcal {D}} \left[ \frac {\mathcal {D} _ {w} (s , a)}{\mathcal {D} (s , a)} \log \frac {\mathcal {D} _ {w} (s , a)}{\mathcal {D} (s , a)} \right] \\ = \mathbb {E} _ {(s, a) \sim \mathcal {D}} \left[ w (s, a) \log w (s, a) \right]. \\ \end{array}
$$
Thus, we can rewrite the training objective (Equation 10) to the follows:
$$
\max  _ {\phi , \psi} \mathbb {E} _ {(s, a, s ^ {\prime}) \sim \mathcal {D}} \left[ \underbrace {w _ {\phi , \psi} (s , a) r (s , a)} _ {\text {R e t u r n}} - \lambda_ {F} \underbrace {\left(w _ {\phi} \left(s ^ {\prime}\right) - w _ {\phi , \psi} (s , a)\right) ^ {2}} _ {\text {B e l l m a n f l o w c o n s e r v a t i o n p e n a l t y}} - \underbrace {\lambda_ {K} w _ {\phi , \psi} (s , a) \log w _ {\phi , \psi} (s , a)} _ {\text {K L r e g u l a r i z a t i o n}} \right]. \tag {24}
$$
The objective in Equation 24 can be optimized by minimizing the loss function  $L(\phi, \psi)$  using stochastic gradient descent shown below:
$$
L (\phi , \psi) := L _ {R} (\phi , \psi) + \lambda_ {F} L _ {F} (\phi , \psi) + \lambda_ {K} L _ {K} (\phi , \psi) \tag {25}
$$
$$
L _ {R} (\phi , \psi) := - \sum_ {i = 1} ^ {B} \bar {w} _ {\phi , \psi} \left(s _ {i}, a _ {i}\right) r \left(s _ {i}, a _ {i}\right) \tag {26}
$$
$$
L _ {F} (\phi , \psi) := \frac {1}{B} \sum_ {i = 1} ^ {B} \left(w _ {\phi , \psi} \left(s _ {i} ^ {\prime}\right) - w _ {\phi , \psi} \left(s _ {i}, a _ {i}\right)\right) ^ {2}. \tag {27}
$$
$$
L _ {K} (\phi , \psi) := \sum_ {i = 1} ^ {B} \bar {w} _ {\phi , \psi} (s, a) \log \bar {w} _ {\phi , \psi} (s, a), \tag {28}
$$
where  $B$  denote batch size and  $(s_i, a_i)$  denotes the  $i^{\text{th}}$  state-action pair in the batch.  $\bar{w}_{\phi, \psi}(s_i, a_i)$  denotes the batch normalized weights predictions [31] that is defined as follows:
$$
\bar {w} _ {\phi , \psi} \left(s _ {i}, a _ {i}\right) := \frac {w _ {\phi , \psi} \left(s _ {i} , a _ {i}\right)}{\sum_ {j = 1} ^ {B} w _ {\phi , \psi} \left(s _ {j} , a _ {j}\right)}, \tag {29}
$$
where  $w_{\phi, \psi}(s_i, a_i)$  denotes the importance weights predictions from the network  $\phi$  and  $\psi$  (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.,  $\mathbb{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  $\phi, \psi$ , 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  $\phi$  and  $\psi$  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  $\phi$  and  $\psi$  for one gradient step as well.
We conducted hyperparameter search in the datasets with diverse trajectories. For AW with CQL we searched temperature  $\eta$ : 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  $\eta$ : 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  $\eta$  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  $\lambda_{K} \in \{0.2, 1.0\}$  and  $\lambda_{F} \in \{0.1, 1.0, 5.0\}$ . We use the best found hyperparameter by the time we started the large scale experiments:  $(\lambda_{K}, \lambda_{F}) = (0.2, 0.1)$  for CQL and  $(\lambda_{K}, \lambda_{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 - \sigma\%$  of trajectories from the random-v2 dataset (low-performing) and  $\sigma\%$  of trajectories from the medium-v2 or expert-v2 dataset (high-performing) for each locomotion environment in the D4RL benchmark. For instance, a dataset that combines  $1 - \sigma\%$  of random and  $\sigma\%$  of medium trajectories is denoted as random-medium- $\sigma\%$ . We evaluate our method and the baselines on these imbalanced datasets across four  $\sigma \in \{1,5,10,50\}$ , four environments. We construct 32 of this type of datasets in all the combinations of {ant, halfcheetah, hopper, walker2d}  $\times$  {random-medium, random-expert}  $\times$  {  $\sigma \in \{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  $\sigma = 1$  and  $\sigma = 5$ , we only evaluate all the methods in  $\sigma = 10$ . Note that the results in  $\sigma = 50$  is not different from the results in larger version of mixed dataset with  $\sigma = 50$  since the amount of high-performing trajectories is still sufficient when  $\sigma = 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 - \sigma\%$  of trajectory segments from the random-v2 dataset (representing low-performing policies) and  $\sigma\%$  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  $\tau_{i}$ , a trajectory segment is selected, ranging in length from 10 to 50 timesteps within  $\tau_{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.