SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Kai Yan Alexander G. Schwing Yu-Xiong Wang
{kaiyan3, aschwing, yxw}@illinois.edu
https://github.com/KaiYan289/TAILO
Offline imitation from observations aims to solve MDPs where only task-specific expert states and task-agnostic non-expert state-action pairs are available. Offline imitation is useful in real-world scenarios where arbitrary interactions are costly and expert actions are unavailable. The state-of-the-art 'Distribution Correction Estimation' (DICE) methods minimize divergence of state occupancy between expert and learner policies and retrieve a policy with weighted behavior cloning; however, their results are unstable when learning from incomplete trajectories, due to a non-robust optimization in the dual domain. To address the issue, in this paper, we propose Trajectory-Aware Imitation Learning from Observations (TAILO). TAILO uses a discounted sum along the future trajectory as the weight for weighted behavior cloning. The terms for the sum are scaled by the output of a discriminator, which aims to identify expert states. Despite simplicity, TAILO works well if there exist trajectories or segments of expert behavior in the task-agnostic data, a common assumption in prior work. In experiments across multiple testbeds, we find TAILO to be more robust and effective, particularly with incomplete trajectories.
In recent years, Reinforcement Learning (RL) has been remarkably successful on a variety of tasks, from games [57] and robot manipulation [66] to recommendation systems [9] and large language model fine-tuning [53]. However, RL often also suffers from the need for extensive interaction with the environment and missing compelling rewards in real-life applications [37].
To address this, Imitation Learning (IL), where an agent learns from demonstrations, is gaining popularity recently [20, 22, 27]. Offline imitation learning, such as behavior cloning (BC) [51], allows the agent to learn from existing experience without environment interaction and reward label, which is useful when wrong actions are costly. However, similar to RL, IL also suffers when limited data is available [52] – which is common as demonstrations of the target task need to be collected every time a new task is addressed. In this work, we consider a special but widely studied [31, 36, 60] case of expert data shortage in offline IL, i.e., offline Learning from Observations (LfO) [71]. In LfO the task-specific data only consists of a few expert trajectories, key frames, or even just the goal (the latter is also known as example-based IL [13]), and the dynamics must be mostly learned from task-agnostic data, i.e., demonstration from data not necessarily directly related to the target task. For example, sometimes the agent must learn from experts with different embodiment [41], where expert actions are not applicable.
In the field of offline LfO, researchers have explored action pseudo-labeling [33, 60], inverse RL [31, 61, 73], and similarity-based reward labeling [8, 55]; for example-based IL, the benchmark is
(a) Problem Settings
(b) Discriminator  $R(s)$
(c) Weighted Behavior Cloning
Figure 1: An illustration of our method, TAILO. Different trajectories are illustrated by different styles of arrows. TAILO consists of two parametric steps: 1) train a discriminator which gives high  $R(s)$  for near-expert states and low  $R(s)$  for non-expert states, as shown in panel b); 2) conduct weighted behavior cloning with weights calculated from  $R(s)$  along the trajectory, as shown in panel c). High transparency indicates a small weight for the state and its corresponding action.
RCE [13], which uses RL with classifier-labeled reward. Recently, Distribution Correction Estimation (DICE) methods, LobsDICE [23] and SMODICE [41], achieve the state of the art for both offline LfO and example-based IL. Both methods minimize the state visitation frequency (occupancy) divergence between expert and learner policies, and conduct a convex optimization in the dual space.
However, DICE methods are neither robust to incomplete trajectories in the task-agnostic / task-specific data [72], nor do they excel if the task-agnostic data contains a very small portion of expert trajectories or only segments [56]. These are inherent shortcomings of the DICE-based methods, as they are indirect: they first perform optimization in a dual domain (equivalent to finding the value function in RL), and then recover the policy by weighted behavior cloning. Importantly, Kullback-Leibler(KL)-based optimization of dual variables requires complete trajectories in the task-agnostic data to balance all terms in the objective. Note, SMODICE with  $\chi^2$ -based optimization is also theoretically problematic (see Appendix C) and empirically struggles on testbeds [41].
To overcome the shortcomings listed above, we propose a simple but effective method for imitation learning from observations and examples, Trajectory-Aware Imitation Learning from Observations (TAILO). We leverage the common assumption that there exist trajectories or long segments that are near-optimal to the task of interest in the task-agnostic data. This assumption is the basis of skill-based learning [20, 46, 47], and the benchmarks of many recent works satisfy this assumption [36, 40, 41, 56]; one real-life example fulfilling this assumption is robotics: the robot often utilizes overlapping skills such as moving the robotic arm and grabbing items from other tasks to complete the current task. Based on this assumption, we discriminate/identify which state-action pairs could be taken by the expert, and assign large weights for trajectory segments leading to those segments in the downstream Weighted Behavior Cloning (WBC). This is a simple way to make the learned policy trajectory-aware. The method only consists of two parametric steps: 1) train a discriminator using positive-unlabeled (PU) learning, and 2) use Weighted Behavior Cloning (WBC) on all state-action pairs in the task-agnostic data with the weights of WBC being a discounted sum over thresholded scores given by the discriminator. Note, the discounted sum propagates large scores to trajectory segments in the past far from expert states, if they lead to expert trajectory segments eventually. Meanwhile, as the task-agnostic data contains both expert and non-expert demonstrations, Positive-Unlabeled (PU) learning is better than plain binary classification. Fig. 1 summarizes our algorithm. We found this simple solution to be surprisingly effective across multiple testbeds, especially if the task-agnostic data contains incomplete trajectories. In this latter case, baselines struggle or even diverge. Moreover, we find our method to also improve if the task-agnostic data contains only a small portion of expert trajectories.
We summarize our contributions as follows: 1) We carefully analyzed the state-of-the-art DICE methods in offline LfO, pointing out their limitations both empirically and theoretically (see Appendix C for details); 2) We propose a simple yet effective solution to offline imitation learning from observations; and 3) We empirically show that this simple method is robust and works better than the state of the art on a variety of settings, including incomplete trajectories, few expert trajectories in the task-agnostic dataset, example-based IL and learning from mismatching dynamics.
Markov Decision Process. A Markov Decision Process (MDP) is a well-established framework for sequential decision-making problems. An MDP is defined by the tuple  $(S, A, T, r, \gamma)$ , where  $S$  is the state space and  $A$  is the action space. For every timestep  $t$  of the Markov process, a state  $s_t \in S$  is given, and an action  $a_t \in A$  is chosen by the agent according to its policy  $\pi(a_t | s_t) \in \Delta(A)$ , where  $\Delta(A)$  is the probability simplex over  $A$ . Upon executing the action  $a_t$ , the MDP will transit to a new state  $s_{t+1} \in S$  according to the transition probability  $T(s_{t+1} | s_t, a_t)$  while the agent receives reward  $r(s_t, a_t) \in \mathbb{R}$ . The goal of the agent is to maximize the discounted reward  $\sum_{t} \gamma^t r(s_t, a_t)$  with discount factor  $\gamma \in [0, 1]$  over a complete run, which is called an episode. The state(-action pairs) collected through the run are called a state(-action) trajectory  $\tau$ . Trajectory segment in this work is defined as a continuous subsequence of a trajectory  $\tau$ . The state visitation frequency (state occupancy) of a policy  $\pi$  is denoted as  $d^\pi(s) = (1 - \gamma) \sum_{t} \gamma^t \operatorname*{Pr}(s_t = s)$ . See Appendix A for a detailed definition of state occupancy and other occupancies.
Positive-Unlabeled Learning. Positive-Unlabeled learning [12] addresses the problem of binary classification with feature  $x \in \mathbb{R}^n$  and label  $y \in \{0, 1\}$  when only the positive dataset  $D_P$  and the unlabeled dataset  $D_U$  are known. Our solution leverages positive prior  $\eta_p = \operatorname*{Pr}(y = 1)$  and negative prior  $\eta_n = \operatorname*{Pr}(y = 0)$ , which are unknown and treated as a hyperparameter.
Offline Imitation from Observations / Examples. Offline imitation learning from observations requires the agent to learn a good policy from two sources of data: one is the task-specific dataset  $D_{\mathrm{TS}}$ , which contains state trajectories  $\tau_{\mathrm{TS}} = \{s_1, s_2, \ldots, s_{n_1}\}$  from the expert that directly addresses the task of interest; the other is the task-agnostic dataset  $D_{\mathrm{TA}}$ , which contains state-action trajectories  $\tau_{\mathrm{TA}} = \{(s_1, a_1), (s_2, a_2), \ldots, (s_{n_2}, a_{n_2})\}$  of unknown optimality to the task of interest. Note that the task-specific trajectory  $\tau_{\mathrm{TS}}$  can be incomplete; specifically, if only the last state exists as an example of success, then it is called imitation from examples [13, 41].
The state of the art methods in this field are SMODICE [41] and LobsDICE [23]. SMODICE minimizes the divergence between the state occupancy from task-specific data  $d^{\mathrm{TS}}$  and the learner policy  $\pi$ 's occupancy  $d^{\pi}$ ; for example, when using a KL-divergence as the metric, the objective is
$$
\min  _ {\pi} \mathrm {K L} \left(d ^ {\pi} (s) \| d ^ {\mathrm {T S}} (s)\right), \text {s . t .} \pi \text {i s a f e a s b l e p o l i c y}. \tag {1}
$$
However, since the task-agnostic dataset is the only source of correspondence between state and action, the state occupancy of the task-agnostic dataset  $d^{\mathrm{TA}}$  must be introduced. With some derivations and relaxations, the objective is rewritten as
$$
\max  _ {\pi} \mathbb {E} _ {s \sim d ^ {\pi}} \log \frac {d ^ {\mathrm {T S}} (s)}{d ^ {\mathrm {T A}} (s)} - \operatorname {K L} \left(d ^ {\pi} (s) \| d ^ {\mathrm {T A}} (s)\right), \text {s . t .} \pi \text {i s a f e a s i b l e p o l i c y}. \tag {2}
$$
Here, the first term  $R(s) = \log \frac{d^{\mathrm{TS}}(s)}{d^{\mathrm{TA}}(s)}$  is an indicator for the importance of the state; high  $R(s)$  means that the expert often visits state  $s$ , and  $s$  is a desirable state. Such  $R(s)$  can be trained by a discriminator  $c(s)$ : a positive dataset  $D_{\mathrm{TS}}$  (label 1) and a negative dataset  $D_{\mathrm{TA}}$  (label 0) are used to find an 'optimal' discriminator  $c = c^{*}$ . Given this discriminator, we have  $R(s) = \log \frac{c^{*}(s)}{1 - c^{*}(s)}$ . SMODICE then converts the constrained Eq. (2) to its unconstrained Lagrangian dual form with dual variable  $V(s)$ , and optimizes the following objective (assuming KL-divergence as the metric):
$$
\min  _ {V} (1 - \gamma) \mathbb {E} _ {s \sim p _ {0}} [ V (s) ] + \log \mathbb {E} _ {(s, a, s ^ {\prime}) \sim D _ {\mathrm {T A}}} \exp [ R (s) + \gamma V (s ^ {\prime}) - V (s) ], \tag {3}
$$
where  $\gamma$  is the discount factor and  $p_0$  is the distribution of the initial state in the MDP. In this formulation,  $R(s)$  can be regarded as the reward function, while  $V(s)$  is the value function. The whole objective is an optimization of a convex function with respect to the Bellman residual. With  $V(s)$  learned, the policy is retrieved via weighted behavior cloning where the coefficient is determined by  $V(s)$ . LobsDICE is in spirit similar, but considers the occupancy of adjacent state pairs instead of a single state.
3.1 Motivation and Overview
As mentioned in Sec. 2, the DICE methods for offline LfO discussed above consist of three parts: reward generation, optimization of the value function  $V(s)$ , and weighted behavior cloning. Such a pipeline can be unstable for two reasons. First, the method is indirect, as the weight for behavior cloning depends on the learned value function  $V(s)$ , which could be inaccurate if the task-agnostic dataset is noisy or is not very related to the task of interest. This is aggravated by the fact that  $V(s')$  as a 1-sample estimation of  $\mathbb{E}_{s' \sim p(s'|s,a)} V(s')$  and logsumexp are used in the objective, which further destabilizes learning. Second, for KL-based metrics, if no state appears twice in the task-agnostic data, which is common for high-dimensional environments, the derivative of the objective with respect to  $V(s)$  is determined by the initial state term, the current state term  $-V(s)$ , and the next state term  $\gamma V(s')$ . Thus, if a non-initial state  $s'$  is missing from the trajectory, then only  $-\gamma V(s')$  remains, which makes the objective monotonic with respect to the unconstrained  $V(s')$ . Consequently,  $V$  diverges for  $s'$  (see Appendix C in the for a more detailed analysis and visualization in Fig. 7).
To address the stability issue, we develop Trajectory-Aware Imitation Learning from Observations (TAILO). TAILO also seeks to find an approximate reward  $R(s)$  by leveraging the discriminator  $c(s)$ , which is empirically a good metric for optimality of the state. However, compared to DICE methods which determine the weight of behavior cloning via a dual program, we adopt a much simpler idea: find 'good' trajectory segments using the discounted sum of future  $R(s)$  along the trajectory following state  $s$ , and encourage the agent to follow them in the downstream weighted BC. To do so, we assign a much larger weight, a thresholding result of the discounted sum, to the state-action pairs for 'good' segments. Meanwhile, small weights on other segments serve as a regularizer of pessimism [25]. Such a method is robust to missing steps in the trajectory. Empirically we find that the weight need not be very accurate for TAILO to succeed. In the remainder of the section, we discuss the two steps of TAILO: 1) training a discriminator to obtain  $R(s)$  (Sec. 3.2), and 2) thresholding over discounted sums of  $R(s)$  along the trajectory (Sec. 3.3).
3.2 Positive-Unlabeled Discriminator
Following DICE,  $R(s)$  is used as a metric for state optimality, and is obtained by training a discriminator  $c(s)$ . However, different from DICE which regards all unlabeled data as negative samples, we use Positive-Unlabeled (PU) learning to train  $c(s)$ , since there often are some expert trajectories or segments of useful trajectories in the task-agnostic dataset. Consequently, we treat the task-agnostic dataset as an unlabeled dataset with both positive samples (expert of the task of interest) and varied negative samples (non-expert).
Our training of  $c(s)$  consists of two steps: 1) training another discriminator  $c'(s)$  that identifies safe negative samples, and 2) formal training of  $c(s)$ . In the first step, to alleviate the issue of treating positive samples from  $D_{\mathrm{TA}}$  as negatives, we use a debiasing objective [29] for the training of  $c'(s)$  (see Appendix A for a detailed derivation):
$$
\min  _ {c ^ {\prime} (s)} L _ {1} \left(c ^ {\prime}\right) = \min  _ {c ^ {\prime} (s)} - \left[ \eta_ {p} \mathbb {E} _ {s \sim D _ {\mathrm {T S}}} \log c ^ {\prime} (s) + \max  \left(0, \mathbb {E} _ {s \sim D _ {\mathrm {T A}}} \log \left(1 - c ^ {\prime} (s)\right) - \eta_ {p} \mathbb {E} _ {s \sim D _ {\mathrm {T S}}} \log \left(1 - c ^ {\prime} (s)\right)\right) \right], \tag {4}
$$
where the bias comes from viewing unlabeled samples as negative samples. Here, positive class prior  $\eta_p$  is a hyperparameter; in experiments, we find results to not be sensitive to this hyperparameter.
In the second step, after  $c^{\prime}(s)$  is trained,  $R^{\prime}(s) = \log \frac{c^{\prime}(s)}{1 - c^{\prime}(s)}$  is calculated for each state in the task-agnostic data, and the states in the (possibly incomplete) trajectories with the least  $\beta_{1} \in (0,1)$  portion of average  $R^{\prime}(s)$  are identified as "safe" negative samples. Note, we do not identify "safe" positive samples, because a trajectory that only has a segment useful for the task of interest might not have the highest average  $R^{\prime}(s)$  due to its irrelevant part; however, an irrelevant trajectory will probably have the lowest  $R^{\prime}(s)$  throughout the whole trajectory, and thus can be identified as a "safe" negative sample. We collect such samples to form a new dataset  $D_{\mathrm{safeTA}}$ .
Finally, the formal training of  $c(s)$  uses states from  $D_{\mathrm{TS}}$  as positive samples and  $D_{\mathrm{safe~TA}}$  as "safe" negative samples. The training objective of  $c(s)$  is a combination of debiasing objective and standard cross entropy loss for binary classification, controlled by hyperparameter  $\beta_{2}$ . Specifically, we use
$$
\min  _ {c (s)} \beta_ {2} L _ {2} (c) + (1 - \beta_ {2}) L _ {3} (c), \text {w h e r e}
$$
$$
L _ {2} (c) = \min  _ {c (s)} - [ \eta_ {p} \mathbb {E} _ {s \sim D _ {\mathrm {T S}}} \log c (s) + \max  (0, \mathbb {E} _ {s \sim D _ {\mathrm {s a f e T A}}} \log (1 - c (s)) - \eta_ {p} \mathbb {E} _ {s \sim D _ {\mathrm {T S}}} \log (1 - c (s))) ],
$$
$$
L _ {3} (c) = \mathbb {E} _ {s \sim D _ {\mathrm {T S}}} \log c (s) + \mathbb {E} _ {s \sim D _ {\text {s a f e T A}}} \log (1 - c (s)). \tag {5}
$$
In this work, we only consider  $\beta_{2} \in \{0,1\}$ ; empirically, we found that  $\beta_{2} = 0$  is better if the agent's embodiments across  $D_{\mathrm{TS}}$  and  $D_{\mathrm{TA}}$  are the same, and  $\beta_{2} = 1$  is better if the embodiments differ. This is because a debiasing objective assumes positive samples still exist in the safe negatives, which pushes the classification margin further from samples in  $D_{\mathrm{TS}}$ , and has a larger probability to classify expert segments in  $D_{\mathrm{TA}}$  as positive.
3.3 Trajectory-Aware Thresholding
With the discriminator  $c(s)$  trained and  $R(s) = \log \frac{c(s)}{1 - c(s)}$  obtained, we can identify the most useful trajectory segments for our task. To do this, we use a simple thresholding which makes the weight employed in behavior cloning trajectory-aware. Formally, the weight  $W(s_{i},a_{i})$  for weighted behavior cloning is calculated as
$$
W \left(s _ {i}, a _ {i}\right) = \sum_ {j = 0} ^ {\infty} \gamma^ {j} \exp \left(\alpha R \left(s _ {i + j}\right)\right), \tag {6}
$$
where  $(s_i, a_i)$  is the  $i$ -th step in a trajectory, and  $\alpha$  is a hyperparameter that controls the strength of thresholding;  $\alpha$  balances the tradeoff between excluding non-expert and including expert-trajectories. For an  $i + j$  which exceeds the length of the trajectory, we set  $s_{i + j}$  to be the last state of the (possibly incomplete) known trajectory, as the final state is of significant importance in many applications [18]. This design also allows us to conveniently address the example-based offline IL problem, where the final state is important. With the weights determined, we finally conduct a weighted behavior cloning with the objective  $\max_{\pi} \mathbb{E}_{(s,a) \sim D_{\mathrm{TA}}} W(s,a) \log \pi(a|s)$ , where  $\pi(a|s)$  is the desired policy.
In this section, we evaluate TAULO on five different, challenging tasks across multiple mujoco testbeds. More specifically, we study the following two questions: 1) Is the algorithm indeed robust to incomplete trajectories in either task-agnostic (Sec. 4.1) or task-specific (Sec. 4.2) data, and does it work with little expert data in the task-agnostic dataset (Sec. 4.3)? 2) Can the algorithm also work well in example-based IL (Sec. 4.4) and learn from experts of different dynamics (Sec. 4.5)?
Baselines. We compare TAULO to four baselines: SMODICE [41], LobsDICE [23], ORIL [73], and Behavior Cloning (BC). Since LobsDICE works with state-pair occupancy, it cannot solve example-based IL; thus, we substitute LobsDICE in example-based IL with RCE [13], a state-of-the-art example-based RL method. Unless otherwise specified, we use 3 random seeds per method in each scenario.
Environment Setup. Following SMODICE [41], unless otherwise specified, we test our algorithm on four standard mujoco testbeds from the OpenAI Gym [5], which are the hopper, halfcheetah, ant, and walker2d environment. We use normalized average reward<sup>1</sup> as the main metric, where higher reward indicates better performance; for environments where the final reward is similar, fewer gradient steps in weighted behavior cloning indicates better performance. We report the change of the mean and standard deviation of reward with respect to the number of gradient steps.
Experimental Setup. For all environments, we use  $\alpha = 1.25$ ,  $\beta_{1} = 0.8$ ,  $\eta_{p} = 0.2$ ;  $\beta_{2} = 1$  if  $D_{\mathrm{TS}}$  and  $D_{\mathrm{TA}}$  are generated by different embodiments, and  $\beta_{2} = 0$  otherwise. We use  $\gamma = 0.998$  unless otherwise specified. We use an exactly identical discriminator and policy network as SMODICE: for the discriminator  $c(s)$  and  $c'(s)$ , we use a small Multi-Layer Perceptron (MLP) with two hidden layers, width 256, and tanh activation function. For actor  $\pi$ , we use an MLP with two hidden layers, width 256, and ReLU [2] activation function. For the training of  $c(s)$  and  $c'(s)$ , we use a learning
${}^{1}$  Normalization standard is according to D4RL [16],and identical to SMODICE [41].
Figure 2: Reward curves for offline imitation learning with incomplete trajectories in the task-agnostic dataset, where the x-axis is the number of gradient steps and the y-axis is the normalized reward. The title for each subfigure is in the format of "environment name" + "1/x" (task-agnostic data removed), where  $x \in \{2, 3, 5, 10, 20\}$ . We observe the proposed method to be the most stable.
rate of 0.0003 and a 1-Lipschitz regularizer, run  $10\mathrm{K}$  gradient steps for  $c^{\prime}(s)$ , and run  $40\mathrm{K}$  gradient steps for  $c(s)$ . For weighted BC, we use a learning rate of  $10^{-4}$ , a weight decay of  $10^{-5}$ , and run 1M gradient steps. For discriminator training, we use a batch size of 512; for weighted BC steps, we use a batch size of 8192. Adam optimizer [28] is used for both steps. See Appendix D for more details and Appendix F for a sensitivity analysis regarding the batch size,  $\alpha, \beta_{1}, \beta_{2}, \eta_{p}$ , and  $\gamma$ .
4.1 Learning from Task-Agnostic Dataset with Incomplete Trajectories
Dataset Setup. We modify the standard dataset settings from SMODICE to create our dataset. SMODICE uses offline datasets from D4RL [16], where a single trajectory from the "expert-v2" dataset is used as the task-specific data. The task-agnostic data consists of 200 expert trajectories (200K steps) from the "expert-v2" dataset and 1M steps from the "random-v2" dataset. Based on this, we iterate over the state-action pairs in the task-agnostic dataset, and remove one pair for every  $x$  pairs. In this work, we test  $x \in \{2,3,5,10,20\}$ .