SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Main Results. Fig. 2 shows the result for different methods with incomplete task-agnostic trajectories, where our method outperforms all baselines and remains largely stable despite decrease of  $x$  (i.e., increase of removed data), as the weights for each state-action pair do not change much. In the training process, we often witness SMODICE and LobsDICE to collapse due to diverging value functions (see Appendix C for explanation), which is expected; for runs that abort due to numerical error, we use a reward of 0 for the rest of the gradient steps. Under a few cases, SMODICE with KL-divergence works decently well with larger noises (e.g., Halfcheetah_1/3 and Walker2d_1/3); this is because sometimes the smoothing effect of the neural network mitigates divergence. However, with larger batch size (See Fig. 23 in Appendix F.4.6) and more frequent and uniform updates on each data point, the larger the noise the harder SMODICE fails.
4.2 Learning from Task-Specific Dataset with Incomplete Trajectories
Dataset Setup and Main Results. We use SMODICE's task-agnostic dataset as described in Sec. 4.1. For the task-specific dataset, we only use the first  $x$  steps and the last  $y$
Hopper head90_tail10
Hopper head50 tail50
Hopper head10 tail90
Hopper_head1.tail100
Halfcheetah head100 tail1
Halfcheetah head90_tail10
Halfcheetah head50 tail50
Halfcheetah head10 tail90
Halfcheetah_head1.tail100
Figure 3: Reward curves for offline imitation with incomplete trajectories in the task-specific dataset. The title format for each subfigure is "environment name" + "head  $x$ " + "tail  $y$ " (states remained), where  $(x, y) \in \{(100, 1), (90, 10), (50, 50), (10, 90), (1, 100)\}$ . We observe the proposed method to be the most stable.
Ant head100 tail1
Ant_head90_tail10
Ant head50 tail50
Ant_head10.tail90
Ant_head1.tail100
Walker2d head90_tail10
Walker2d head50 tail50
Walker2d_head10.tail90
Walker2d_head1.tail100
steps in the expert trajectory, and discard the remainder. In this work, we test  $(x,y) \in \{(1,100), (10,90), (50,50), (90,10), (100,1)\}$ . Fig. 3 shows the result with incomplete task-specific trajectories, where our method outperforms all baselines and often achieves results similar to those obtained when using the entire task-specific dataset. In contrast, SMODICE and LobsDICE are expectedly unstable in this setting.
4.3 Standard Offline Imitation Learning from Observations
Environment Setup. In addition to the four standard mujoco environments specified above, we also test on two more challenging environments: the Franka kitchen environment and the antimaze environment from D4RL [16]. In the former, the agent needs to control a 9-DoF robot arm to complete a sequence of item manipulation subtasks, such as moving the kettle or opening the microwave; in the latter, the agent needs to control a robot ant to crawl through a U-shaped maze and get to a particular location. As the kitchen environment requires less steps to finish, we use  $\gamma = 0.98$  instead of 0.998.
Dataset Setup. As existing methods already solve the four mujoco environments with SMODICE's dataset settings in Sec. 4.1 quite well (see Appendix F for result), we test a more difficult setting to demonstrate TAILO's ability to work well with few expert trajectories in the task-agnostic dataset. More specifically, we use the same task-specific dataset, but only use 40 instead of 200 expert trajectories from the "expert-v2" dataset to mix with the 1M "random-v2" steps data and form the task-agnostic dataset. For the more challenging kitchen and antimaze environment, we use the identical dataset as SMODICE, where a single trajectory is used as the task-specific dataset. The task-agnostic dataset for the kitchen environment consists of expert trajectories completing different subtasks (both relevant and irrelevant) in different orders, and the task-agnostic dataset for antimaze consists of data with varied optimality. See Appendix D for details.
Main Results. Fig. 4 shows the result for different methods in standard offline imitation learning from observations. We find our method to outperform all baselines on hopper, halfcheetah, ant and walker2d. Results are comparable to the best baseline on kitchen and antmaze. In the experiment, we found ORIL to often diverge, and the performance of SMODICE varies greatly depending on the  $f$ -divergence: on hopper, halfcheetah and walker2d, KL-divergence is much better, while  $\chi^2$
Figure 4: Reward curves for offline imitation learning from observations with few expert trajectories in the task-agnostic dataset. In all six environments, our method either outperforms or is comparable to the baselines.
Figure 5: Reward curves for offline imitation learning from examples.
divergence is better on ant and kitchen. LobsDICE is marginally better than SMODICE, as the former considers state-pair occupancy instead of single state occupancy, which is more informative. Also worth noting: none of the methods exceeds a normalized reward of 60 in the kitchen environment; this is because the SMODICE experiment uses expert trajectories that are only expert for the first 2 out of all 4 subtasks. See Sec. F.3 in the Appendix for a detailed discussion.
4.4 Learning from Examples
Environment Setup. Following SMODICE, we test example-based offline imitation in three different testbeds: pointmaze, kitchen and antmaze. In the mujoco-based [59] pointmaze environment the agent needs to control a pointmass on a 2D plane to navigate to a particular direction. For the kitchen environment, we test two different settings where the agent is given successful examples of moving the kettle and opening the microwave respectively (denoted as "kitchen-kettle" and "kitchen-microwave"). As pointmaze requires less steps to finish, we use  $\gamma = 0.98$  instead of 0.998.
Dataset Setup and Main Results. Following SMODICE, we use a small set of success examples:  $\leq 200$  states for antmaze and pointmaze and 500 states for kitchen (see Appendix D for details) are given as the task-specific dataset. For pointmaze, the task-agnostic data contains 60K steps, generated by a script and distributed evenly along four directions (i.e.,  $25\%$  expert data); for other testbeds, the task-agnostic data is identical to SMODICE described in Sec. 4.1. Fig. 5 shows the results for offline imitation learning from examples on all environments tested in SMODICE; TAILO is marginally better than the baselines on pointmaze, antmaze, and kitchen-microwave, and is comparable (all close to perfect) on kitchen-kettle.
4.5 Learning from Mismatched Dynamics
Environment and Dataset Setup. Following SMODICE, we test our method on three environments: antmaze, halfcheetah, and ant. We use task-specific data from an expert with different dynamics (e.g., ant with a leg crippled; see Sec. E for details). Task-agnostic data follows SMODICE in Sec. 4.1.
Main Results. Fig. 6 shows the results for offline imitation learning from mismatched dynamics, where our method is the best in all three testbeds. Among the three environments, halfcheetah is the most difficult, as the state space for halfcheetah with shorter torso is unreachable by a normal
Figure 6: Reward curves for offline imitation learning from mismatched dynamics.
halfcheetah, i.e., the assumption that  $d^{\mathrm{TA}}(s) > 0$  wherever  $d^{\mathrm{TS}}(s) > 0$  in SMODICE and LobsDICE does not hold; in such a setting, our method is much more robust than DICE methods.
Offline Imitation Learning and Distribution Correction Estimation (DICE). Offline imitation learning aims to learn a policy only from data without interaction with the environment. This is useful where immature actions are costly, e.g., in a dangerous factory. The simplest solution for offline imitation learning is plain Behavior Cloning (BC) [51]. Many more methods have been proposed recently, such as BCO [60] and VMSR [33] (which pseudolabel actions), offline extensions of GAIL [3, 31, 61, 73], and similarity-based reward labeling [8, 35, 55, 64]. Currently, the state-of-the-art method for offline imitation learning is DDistribution Correction Estimation (DICE) [23, 27, 32, 34, 41], which minimizes the discrepancy between an expert's and a learner's state, state-action, or state-pair occupancy. Our method is inspired by DICE, but is much simpler and more effective. More recently, two works unify offline imitation learning with offline RL, which are offline-RL-based MAHALO [36] and DICE-based ReCOIL [56]. Different from SMODICE and LobsDICE, ReCOIL minimizes the divergence between learner and expert data mixed with non-expert data respectively, which removes the data coverage assumption. However, ReCOIL faces the problem discussed in Sec. 3.1 when dealing with incomplete trajectories, and MAHALO is based on state-pairs similar to LobsDICE, which cannot solve IL with incomplete trajectories or example-based IL like our TAULO.
Learning from Observations and Examples. Learning from observations (LfO) [60] requires the agent to learn from a task-specific dataset without expert action, which is useful when learning from videos [44] or experts with different embodiments [55], as the expert action is either unavailable or not applicable. Learning from examples is an extreme case of LfO where only the final goal is given [13]. There are three major directions: 1) pseudolabeling of actions which builds an inverse dynamic model and predicts the missing action [33, 60]; 2) occupancy divergence minimization with either DICE [23, 27, 34, 41, 71] or inverse-RL style iterative update [61, 67, 73]; and 3) RL/planning with reward assignment based on state similarity (often in visual imitation) [9, 55, 64]. Our proposed TAILO solves LfO with a simple solution different from existing ones.
Discriminator as Reward. The idea of training a discriminator to provide rewards for states is widely used in IL, including inverse RL methods [15, 22, 31, 73], DICE methods [23, 27, 41], and methods such as 2IWIL [65] and DWBC [68]. In this work, we propose a simple but explicit way to take the trajectory context into account, using the output from a discriminator as reward, which differs from prior works.
Positive-Unlabeled Learning. Positive-Unlabeled (PU) [11, 12, 29, 50] learning aims to solve binary classification tasks where only positive and unlabeled data are available. It is widely used in data retrieval [54], outlier detection [38], recommendation [7] and control tasks [67]. In this work, we utilize two PU learning achievements: the skill of identifying "safe" negative samples [39] and debiasing [11, 29]. The closest RL work to our use of PU learning is ORIL [73], which also uses positive unlabeled learning to train a discriminator for reward estimation. However, there are three key differences between our method and ORIL: we define  $R(s)$  differently for better thresholding, use different techniques for PU learning to prevent overfitting, and, importantly, the removal of value function learning. See Appendix F.6 for an ablation to demonstrate efficacy of these changes.
Reward-Weighted Regression(RWR) [48] and Advantage-Weighted Regression(AWR) [45]. The idea in our Eq. (6) of weighted behavior cloning with weights being (often exponentiated) discounted return has been widely used in the RL community [1, 45, 63], and our objective resembles that of
RWR/AWR. However, our work differs from RWR/AWR inspired works [30, 43, 49] in the following aspects:
- The objective for both AWR and RWR are built upon the related payoff procedure [10, 21], which introduces an Expectation-Maximization (EM) procedure for RL. However, for offline IL in our case, iteration between E-step and M-step are infeasible. There are two workarounds for this: importance sampling and naively using one iteration. However, the former is known to be non-robust [42] and the latter, MARWIL [62], struggles in our testbeds (see Appendix F.2).
- We use neither an adaptive reward scaling term [48] in the EM framework nor parametric value estimation [45, 62], which are both widely utilized by RWR/AWR inspired works. However, the former is not guaranteed to preserve an optimal policy and thus is not an advantage [58], while the latter struggles in our setting where learning a good value function is hard, as illustrated by the performance of baselines such as DICE methods and MARWIL.
- For all existing RWR/AWR works, the reward labels are assumed to be available, which is different from our case.
We propose TAULO, a simple yet effective solution for offline imitation from observations by training a discriminator using PU learning, applying a score to each state-action pair, and then conducting a weighted behavior cloning with a discounted sum over thresholded scores along the future trajectory to obtain the weight. We found our method to improve upon state-of-the-art baselines, especially when the trajectories are incomplete or the expert trajectories in the task-agnostic data are few.
Societal Impact. Our work addresses sequential decision-making tasks from expert observations with fewer related data, which makes data-driven automation more applicable. This, however, could also lead to negative impacts on society by impacting jobs.
Limitations and Future Directions. Our method relies on the assumption that there exist expert trajectories or at least segments in the task-agnostic data. While this is reasonable for many applications, like all prior work with proprioceptive states, it is limited in generalizability. Thus, one promising future direction is to improve the ability to summarize abstract "skills" from data with better generalizability. Another future direction is learning from video demonstrations, which is a major real-world application for imitation learning from observations. Also, our experiments are based on simulated environments such as D4RL. Thus a gap between our work and real-life progress remains. While we are following the settings of many recent works, such as SMODICE [41], ReCOIL [56] and OTR [40], to bridge the gap using techniques such as sim2real in the robotics community [24] is another very important direction for future work.
Acknowledgements. This work was supported in part by NSF under Grants 2008387, 2045586, 2106825, MRI 1725729, NIFA award 2020-67021-32799, the Jump ARCHES endowment through the Health Care Engineering Systems Center, the National Center for Supercomputing Applications (NCSA) at the University of Illinois at Urbana-Champaign through the NCSA Fellows program, the IBM-Illinois Discovery Accelerator Institute, and the Amazon Research Award.
[1] A. Abdolmaleki, J. T. Springenberg, Y. Tassa, R. Munos, N. Heess, and M. Riedmiller. Maximum a posteriori policy optimisation. *ArXiv:1806.06920*, 2018.
[2] A. F. Agarap. Deep learning using rectified linear units (relu). ArXiv:1803.08375, 2018.
[3] S. Y. Arnob. Off-policy adversarial inverse reinforcement learning. ArXiv:2005.01138, 2020.
[4] M. I. Belghazi, A. Baratin, S. Rajeshwar, S. Ozair, Y. Bengio, A. Courville, and D. Hjelm. Mutual information neural estimation. In ICML, 2018.
[5] G. Brockman, V. Cheung, L. Pettersson, J. Schneider, J. Schulman, J. Tang, and W. Zaremba. Openai gym, 2016.
[6] A. Camacho, I. Gur, M. Moczulski, O. Naschum, and A. Faust. Sparsedice: Imitation learning for temporally sparse data via regularization. In the Unsupervised Reinforcement Learning Workshop in ICML, 2021.
[7] S. Chang, Y. Zhang, J. Tang, D. Yin, Y. Chang, M. A. Hasegawa-Johnson, and T. S. Huang. Positive-unlabeled learning in streaming networks. In KDD, 2016.
[8] A. S. Chen, S. Nair, and C. Finn. Learning generalizable robotic reward functions from "in-the-wild" human videos. ArXiv:2103.16817, 2021.
[9] X. Chen, S. Li, H. Li, S. Jiang, Y. Qi, and L. Song. Generative adversarial user model for reinforcement learning based recommendation system. In ICML, 2019.
[10] P. Dayan and G. E. Hinton. Using expectation-maximization for reinforcement learning. Neural Computation, 1997.
[11] M. C. du Plessis, G. Niu, and M. Sugiyama. Analysis of learning from positive and unlabeled data. In NIPS, 2014.
[12] C. Elkan and K. Noto. Learning classifiers from only positive and unlabeled data. In KDD, 2008.
[13] B. Eysenbach, S. Levine, and R. Salakhutdinov. Replacing rewards with examples: Example-based policy search via recursive classification. In NeurIPS, 2021.
[14] G. Freund, E. Sarafian, and S. Kraus. A coupled flow approach to imitation learning. In ICML, 2023.
[15] J. Fu, K. Luo, and S. Levine. Learning robust rewards with adversarial inverse reinforcement learning. In ICLR, 2018.
[16] J. Fu, A. Kumar, O. Nachum, G. Tucker, and S. Levine. D4rl: Datasets for deep data-driven reinforcement learning. ArXiv:2004.07219, 2020.
[17] S. Fujimoto, H. Hoof, and D. Meger. Addressing function approximation error in actor-critic methods. In ICML, 2018.
[18] S. Ghasemipour, R. Zemel, and S. Gu. A divergence minimization perspective on imitation learning methods. In CoRL, 2019.
[19] A. L. Gibbs and F. E. Su. On choosing and bounding probability metrics. International statistical review, 2002.
[20] K. Hakhamaneshi, R. Zhao, A. Zhan, P. Abbeel, and M. Laskin. Hierarchical few-shot imitation with skill transition models. In ICLR, 2022.
[21] G. E. Hinton. Connectionist learning procedures. 1990.
[22] J. Ho and S. Ermon. Generative adversarial imitation learning. In NIPS, 2016.
[67] D. Xu and M. Denil. Positive-unlabeled reward learning. In CoRL, 2019.
[68] H. Xu, X. Zhan, H. Yin, and H. Qin. Discriminator-weighted offline imitation learning from suboptimal demonstrations. In NeurIPS, 2022.
[69] T. Xu, Z. Li, Y. Yu, and Z.-Q. Luo. On generalization of adversarial imitation learning and beyond. arXiv preprint arXiv:2106.10424, 2021.
[70] T. Yu, G. Thomas, L. Yu, S. Ermon, J. Y. Zou, S. Levine, C. Finn, and T. Ma. Mopo: Model-based offline policy optimization. In NeurIPS, 2020.
[71] Z. Zhu, K. Lin, B. Dai, and J. Zhou. Off-policy imitation learning from observations. In NeurIPS, 2020.
[72] L. Ziniu, X. Tian, Y. Yang, and L. Zhi-Quan. Rethinking valuedice - does it really improve performance? In ICLR Blog Track, 2022.
[73] K. Zolna, A. Novikov, K. Konyushkova, C. Gulcehre, Z. Wang, Y. Aytar, M. Denil, N. de Freitas, and S. E. Reed. Offline learning from demonstrations and unlabeled experience. In *Offline Reinforcement Learning Workshop at NeurIPS*, 2020.
Appendix: A Simple Solution for Offline Imitation from Observations and Examples with Possibly Incomplete Trajectories
The Appendix is organized as follows: first, we summarize our key findings. Then, in Sec. A, we provide a more rigorous introduction of the key mathematical concepts, and in Sec. B we provide the pseudocode for the training process of TAILO. In Sec. C, we explain why DICE methods struggle with incomplete trajectories in the task-agnostic or task-specific data. We then list additional implementation details of our method and the baselines in Sec. D, include additional experimental settings in Sec. E, and present additional experiment results as well as ablation studies in Sec. F; after that, we report the training time and computational resources utilized by each method in Sec. G. Finally, we examine the licenses of assets for our code in Sec. H. See https://github.com/KaiYan289/TAILO for our code.
The key findings of the additional experimental results are summarized as follows:
- How is our method different from ORIL, and where does the performance gain come from? Our method is different from ORIL in three major aspects: definition of  $R(s)$ , positive-unlabeled learning technique, and policy retrieval. In Sec. F.6, we report that each of the three aspects contributes to the reward increase; among them, policy retrieval is the most important factor, and the positive-unlabeled learning technique is the least.
A Mathematical Concepts
In this section, we rigorously introduce four mathematical concepts in our paper, which are state(-action/-pair) occupancy, debiasing objective for positive-unlabeled learning,  $f$ -divergences, and Fenchel conjugate. The first one is used throughout the paper, the second is used in the training of  $c(s)$ , and the others are used in Sec. C.
State, State-Action, and State-pair Occupancy. Consider an infinite horizon MDP  $(S,A,T,\gamma)$  with initial state distribution  $p_0$ , where the state at the  $t$ -th timestep is  $s_t$  and the action is  $a_t$ . Then, given any fixed policy  $\pi$ , the probability  $\operatorname{Pr}(s_t = s)$  of landing in any state  $s$  for any step  $t$  is determined. Based on this, the state visitation frequency  $d^{\pi}(s)$  with policy  $\pi$ , also known as state occupancy, is defined as  $d^{\pi}(s) = \sum_{t=1}^{\infty} \operatorname{Pr}(s_t = s)$ . Similarly, the state-action frequency is defined as  $d^{\pi}(s,a) = \sum_{t=1}^{\infty} \operatorname{Pr}(s_t = s,a_t = a)$ , and the state-pair frequency is defined as  $d^{\pi}(s,s') = \sum_{t=1}^{\infty} \operatorname{Pr}(s_t = s,s_{t+1} = s')$ . For better readability, with a little abuse of notation, we use the same  $d$  for all three occupancies throughout the paper. In this work, we refer to the occupancy with the learner's policy  $\pi$  as  $d^{\pi}$ , we refer to the occupancy with average policy of the task-agnostic dataset as  $d^{\mathrm{TA}}$ , and to the occupancy with average policy of the task-specific dataset as  $d^{\mathrm{TS}}$ .
Debiasing Objective for Positive-Unlabeled Learning. Consider a binary classification task with feature  $x \in \mathbb{R}^n$  and label  $y \in \{0,1\}$ . Assume we have access to both labeled positive dataset  $D_P$  and labeled negative dataset  $D_N$ , and the class prior, i.e., the probability of having a label when uniformly sampled from the dataset, is  $\eta_p = \operatorname*{Pr}(y = 1)$  for positive labels and  $\eta_n = \operatorname*{Pr}(y = 0)$  for negative samples. Then, with sufficiently many samples, the most commonly used loss function is a cross entropy loss. The average cross entropy loss can be approximated by
$$
\eta_ {p} \mathbb {E} _ {(x, y) \sim D _ {P}} [ - \log \hat {y} ] + \eta_ {n} \mathbb {E} _ {(x, y) \sim D _ {N}} [ - \log (1 - \hat {y}) ], \tag {7}
$$
where  $\hat{y} = c(x)$  is the output of the discriminator.
In positive-unlabeled learning, we only have access to the positive dataset  $D_P$  and unlabeled dataset  $D_U$  as an unlabeled mixture of positive and negative samples. One naive approach is to regard all samples from  $D_U$  as negative samples. While this sometimes works, it falsely uses the loss function for positive samples in  $D_U$ , and thus introduces bias. To avoid this and get better classification results, multiple debiasing objectives [11, 29] have been proposed. Such objectives assume that the unlabeled dataset  $D_U$  consists of  $\eta_p$  portion of positive data and  $\eta_n$  portion of negative data. Thus, we have
$$
\mathbb {E} _ {(x, y) \sim D _ {U}} [ - \log (1 - \hat {y}) ] = \eta_ {p} \mathbb {E} _ {(x, y) \sim D _ {P}} [ - \log (1 - \hat {y}) ] + \eta_ {n} \mathbb {E} _ {(x, y) \sim D _ {N}} [ - \log (1 - \hat {y}) ]. \tag {8}
$$
Correspondingly, Eq. (7) can be approximated by
$$
- \left[ \eta_ {p} \mathbb {E} _ {(x, y) \sim D _ {P}} \log \hat {y} + \mathbb {E} _ {(x, y) \sim D _ {U}} \log (1 - \hat {y}) - \eta_ {p} \mathbb {E} _ {(x, y) \sim D _ {P}} \log (1 - \hat {y}) \right]. \tag {9}
$$
This is the formulation utilized by ORIL [73]. However, as Kiryo et al. [29] pointed out, the red part in Eq. (9) is an approximation for  $\eta_{n}\mathbb{E}_{(x,y)\sim D_{N}}[-\log (1 - \hat{y})]$ , and thus should be no less than 0; violation of such a rule could lead to overfitting (see [29] for details). Therefore, we apply a max operator  $\max (\cdot ,0)$  to the red part. The objective thus becomes
$$
- \left[ \eta_ {p} \mathbb {E} _ {(x, y) \sim D _ {P}} \log \hat {y} + \max  \left(\mathbb {E} _ {(x, y) \sim D _ {U}} \log (1 - \hat {y}) - \eta_ {p} \mathbb {E} _ {(x, y) \sim D _ {P}} \log (1 - \hat {y}), 0\right) \right]. \tag {10}
$$