SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Note that while Kiryo et al. [29] uses an alternative update conditioning on the red term, we found the max operator to be more effective in our case. For our first step of discriminator training, we use Eq. (10), with  $D_{\mathrm{TS}}$  as  $D_P$ ,  $D_{\mathrm{TA}}$  as  $D_U$ , state  $s$  as feature  $x$ , and action  $a$  as label  $y$ . For the second step, we use Eq. (10) again for task-specific data with mismatch dynamics, but this time we use
$D_{\mathrm{TA~safe}}$  as  $D_U$ ; otherwise, we use Eq. (7) with  $D_{\mathrm{TS}}$  as  $D_P$ ,  $D_{\mathrm{TA~safe}}$  as  $D_N$ , state  $s$  as feature  $x$ , and action  $a$  as label  $y$ .
$f$ -divergences. The  $f$ -divergence is the basis of the DICE family of methods [23, 27, 32, 34, 41]. DICE methods minimize an  $f$ -divergence, such as the KL-divergence between the learner's policy and the average policy on the task-specific dataset. For any continuous and convex function  $f$ , any domain  $\mathcal{X}$ , and two probability distributions  $p, q$  on  $\mathcal{X}$ , the  $f$ -divergence between  $p$  and  $q$  is defined as
$$
D _ {f} (p \| q) = \mathbb {E} _ {x \sim q} [ f (\frac {p (x)}{q (x)}) ]. \tag {11}
$$
For example, for the KL-divergence,  $f(x) = x \log x$ , and  $D_{f}(p\| q) = \mathrm{KL}(p\| q) = \mathbb{E}_{x \sim q}[\frac{p(x)}{q(x)} \log \frac{p(x)}{q(x)}] = \mathbb{E}_{x \sim p} \log \frac{p(x)}{q(x)}$ ; for  $\chi^2$ -divergence,  $f(x) = (x - 1)^2$ , and  $D_{f}(p\| q) = \chi^2(p\| q) = \mathbb{E}_{x \sim q}(\frac{p(x)}{q(x)} - 1)^2 = \int_{x \sim \mathcal{X}} \frac{(p(x) - q(x))^2}{q(x)}$ . Note that SMODICE uses  $f(x) = \frac{1}{2}(x - 1)^2$ , which is essentially half  $\chi^2$ -divergence.
Fenchel Conjugate. For any vector space  $\Omega$  with inner product  $\langle \cdot, \cdot \rangle$  and convex and differentiable function  $f: \Omega \to \mathbb{R}$ , the Fenchel conjugate of  $f(x)$ ,  $x \in \Omega$  is defined as
$$
f _ {*} (y) = \max  _ {x \in \Omega} \langle x, y \rangle - f (x). \tag {12}
$$
In SMODICE, the derivation of the  $\chi^2$ -divergence uses  $f(x) = \frac{1}{2} (x - 1)^2$ , and the Fenchel dual is  $f_{*}(y) = \frac{1}{2} (y + 1)^{2}$ . However, such Fenchel dual is obtained with no constraint on  $x$ , while in SMODICE with KL-divergence,  $x$  is constrained on a probability simplex. This extra relaxation could be a factor for its worse performance on some testbeds.
B Algorithm Details
Alg. 1 shows the pseudocode of our method.
C Why DICE Struggles with Incomplete Trajectories?
In our experiments, we empirically find that SMODICE and LobsDICE struggle with either incomplete task-specific or incomplete task-agnostic trajectories. We give an extended explanation in this section.
C.1 Incomplete Task-Specific Trajectory
The phenomenon that DICE struggles with incomplete expert trajectories is first discussed in [72]: the work mentions that subsampled (i.e., incomplete) trajectories artificially "mask" some states and causes the failure of ValueDICE. Recent discussion indicates that such a failure on subsampled task-specific trajectories is closely related to overfitting [6] and a lack of generalizability [69, 72].
C.2 Incomplete Task-Agnostic Trajectory
C.2.1 KL-Based Formulation
The reason that SMODICE [41] with KL-divergence and LobsDICE [23] struggle with incomplete task-agnostic trajectories is more direct: consider the final objective of SMODICE with KL-divergence:
$$
\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 {13}
$$
where  $p_0 \in \Delta(S)$  is the initial state distribution over the state space  $S$ , and "reward function"  $R(s) = \log \frac{d^{\mathrm{TS}}(s)}{d^{\mathrm{TA}}(s)}$  is labeled by a discriminator  $c(s)$  as described in Sec. 2.
Similarly, the final objective of LobsDICE is
$$
\min  _ {V} (1 - \gamma) \mathbb {E} _ {s \sim p _ {0}} V (s) + (1 + \alpha) \log \mathbb {E} _ {(s, a, s ^ {\prime})) \sim D _ {\mathrm {T A}} \exp \left[ \left(\frac {1}{1 + \alpha} \left(R (s, s ^ {\prime}) + \gamma V (s ^ {\prime}) - V (s)\right) \right], \right. \tag {14}
$$
where the "reward function"  $R(s, s')$  is based on a state pair instead of a single state, and  $\alpha > 0$  is a hyperparameter. Note, both methods use a 1-sample estimate for the expectation of the dual variable  $\mathbb{E}_{s'} V(s')$  for future state  $s'$ .
Assume that no two states in  $D_{\mathrm{TA}}$  are exactly the same, which is common in a high-dimensional continuous state space. In such a case, there are only two occurrences of  $V(s)$  for a particular state  $s \in D_{\mathrm{TA}}$  in Eq. (13) and Eq. (14): for initial states, one in the linear term and the other in  $-V(s)$  inside the exp-term; for other states, one in the  $-V(s)$  term and the other in the  $\gamma V(s')$  term inside the exp-term. Consider a trajectory  $\tau = \{(s_1, a_1, s_2), (s_2, a_2, s_3), \ldots, (s_n, a_n, s_{n+1})\} \in D_{\mathrm{TA}}$ . If  $(s_2, a_2, s_3)$  is missing, we have the following ways to make up:
1. Use  $D_{\mathrm{TA}}$  without extra handling. In this case, the only occurrence of  $V(s_3)$  will be  $-V(s_3)$ , as no future state exists for  $s_2$ . As the objective is monotonic (which is not the case in RL) with unconstrained  $V(s_3)$ , the convergence of the algorithm solely relies on smoothing from nearby states in  $D_{\mathrm{TA}}$ . This is also the method that we tested in Sec. F.
2. Let the trajectory contain  $(s_1, a_1, s_3)$ . In this case, the method would learn with a wrong dynamic of the environment because  $s_1$  cannot transit to  $s_3$  by conducting action  $a_1$  as such a ground truth would suggest;
3. Consider  $s_2$  as a terminal state of the trajectory  $\tau$ , where the agent could learn to halt in the middle of an expert trajectory if  $\tau$  is an expert trajectory;
4. Train an inverse dynamic model for pseudolabeling of the action. However, such a method cannot deal with more than one consecutive transition missing.
In conclusion, none of the workarounds described above addresses the concern. Note that the problem described here is also valid for ValueDICE [32] and many follow-up works such as DemoDICE [27] as long as the Donsker-Varadhan representation [4] of the KL-divergence is used. Fig. 7 illustrates this divergence, which empirically verifies our motivation.
(a) Minimum  $V(s)$
(b) Maximum  $V(s)$
Figure 7: Visualization of collapsing patterns of SMODICE-KL in Sec. 4.1 (walker_1/5 is collpasing; walker_1/3 is not). Walker_1/5 stops early due to NaN values in training, which we denote as 0 reward in our paper. It is clearly shown that the reward decrease is closely related to  $V(s)$  divergence.
(c) Normalized Reward
C.2.2  $\chi^2$ -Based Formulation
Empirically, we found that SMODICE with  $\chi^2$ -divergence struggles on several testbeds such as halfcheetah, hopper, and walker2d, which is consistent with the results reported by SMODICE. The performance gap could be due to the following two reasons:
First, violation of Theorem 1 in SMODICE [41]. Theorem 1 in SMODICE reads as follows:
Theorem 1. Given the assumption that  $d^{\mathrm{TA}}(s) > 0$  whenever  $d^{\mathrm{TS}}(s) > 0$ , we have
<sup>2</sup>Because DICE methods utilize the next state  $s'$  for a state-action pair  $(s, a)$ , we write the state-action trajectories as transition  $(s, a, s')$  trajectories.
$$
\mathrm {K L} \left(d ^ {\pi} (s) \| d ^ {\mathrm {T S}} (s)\right) \leq \mathbb {E} _ {s \sim d ^ {\pi}} \left[ \log \left(\frac {d ^ {\mathrm {T A}} (s)}{d ^ {\mathrm {T S}} (s)}\right) \right] + \mathrm {K L} \left(d ^ {\pi} (s, a) \| d ^ {\mathrm {T A}} (s, a)\right), \tag {15}
$$
and furthermore, for any  $f$ -divergence  $D_{f}$  larger than KL,
$$
\operatorname {K L} \left(d ^ {\pi} (s) \| d ^ {\mathrm {T S}} (s)\right) \leq \mathbb {E} _ {s \sim d ^ {\pi}} \left[ \log \left(\frac {d ^ {\mathrm {T A}} (s)}{d ^ {\mathrm {T S}} (s)}\right) \right] + D _ {f} \left(d ^ {\pi} (s, a) \| d ^ {\mathrm {T A}} (s, a)\right). \tag {16}
$$
The right hand side of Eq. (15) is the optimization objective of SMODICE with KL-divergence, and  $\chi^2$ -divergence is introduced via Eq. (16) as an upper bound.
However, SMODICE uses  $f(x) = \frac{1}{2}(x - 1)^2$  instead of  $f(x) = (x - 1)^2$  as the  $\chi^2$ -divergence; i.e., the  $\chi^2$ -divergence is halved in the objective. While the  $\chi^2$ -divergence is an upper bound of the KL-divergence [19], half of the  $\chi^2$ -divergence is not. For example, consider two binomial distributions  $B(1,p)$  and  $B(1,q)$ . The KL-divergence between the two is  $p\log \frac{p}{q} + (1 - p)\log \frac{1 - p}{1 - q}$ , and the  $\chi^2$ -divergence is  $\frac{(p - q)^2}{q} + \frac{(q - p)^2}{1 - q}$ . When  $p = 0.99$  and  $q = 0.9$ , the KL-divergence is  $0.99\log 1.1 + 0.01\log 0.1 \approx 0.0713$ , and the  $\chi^2$ -divergence is  $\frac{0.09^2}{0.9} + \frac{0.09^2}{0.1} = 0.09$ , where half of the  $\chi^2$  divergence is smaller than the KL-divergence. Thus, SMODICE with  $\chi^2$ -divergence is not optimizing an upper bound of  $\mathrm{KL}(d^\pi(s)\|d^{\mathrm{TS}}(s))$  and the performance is not guaranteed. In Fig. 8, we plot the reward curves of using halved and full  $\chi^2$ -divergence under the settings of Sec. F.5; we find that using full  $\chi^2$ -divergence significantly increases performance on the walker2d environment and performs similarly on other environments, though still much worse than our method.
Figure 8: Performance comparison between using halved  $\chi^2$ -divergence (purple) adopted by SMODICE [41] and full  $\chi^2$ -divergence (fuchsia) with experiment settings in Sec. F.5. Full  $\chi^2$ -divergence performs significantly better on the walker2d environment than halved  $\chi^2$ -divergence, though still much worse than our method.
Second, SMODICE with  $\chi^2$ -divergence uses  $f_{*}(y) = \frac{1}{2} (y + 1)^{2}$  as the Fenchel conjugate of  $f(x) = \frac{1}{2} (x - 1)^{2}$ . Such a conjugate is obtained when there is no constraint on  $x$ , i.e.,  $x \in \mathbb{R}$ ; however, such  $x$ , as indicated by the derivation of SMODICE with KL-divergence (Example 1 of Appendix C in SMODICE), should be on the probability simplex. This relaxation could be another reason for the performance drop in SMODICE with  $\chi^2$ -divergence.
C.3 Other Recent Works
We noticed that there are also recent works in the DICE family that try to address the issue of incomplete (i.e., subsampled) trajectories, such as SparseDICE [6] and CFIL [14] for online imitation learning; LobsDICE [23] also discusses learning with subsampled expert trajectories. However, all those works only focus on incomplete task-specific (i.e. expert) trajectories instead of task-agnostic trajectories; also, our method can solve example-based IL and task-specific trajectories with all but one state provided at the beginning, which differs from their settings where the sampling of expert states/state-action pairs is uniform throughout the trajectory.
D Additional Implementation Details
Our code is provided in the supplementary material. We implement our algorithm from scratch, and use the implementation of SMODICE [41] (https://github.com/JasonMa2016/SMODICE) as
the codebase for SMODICE, ORIL [73], and RCE [13]. We obtain the code for LobsDICE [23] from their publicized supplementary material on OpenReview.
Tab. 1 summarizes the unique hyperparameters for our method, and Tab. 2 summarizes the common training paradigms for our method and the baselines. For all baselines, if the hyperparameter values in the paper and the code are different, we record the values from the code. We discuss the influence of batch size and sensitivity of our hyperparameters in Sec. E. Tab. 3 summarizes the hyperparameters specific to other methods. SMODICE, ORIL, and RCE first train a discriminator, and then jointly update actor and critic, while LobsDICE jointly updates all three networks.
E Additional Experimental Settings
In this section, we describe in detail how the environment is setup and how the dataset is generated.
E.1 Offline Imitation Learning from Task-Agnostic Dataset with Incomplete Trajectories
Environment Settings. Sec. 4.1 tests four Mujoco environments: hopper, halfcheetah, ant, and walker2d. The detailed configuration for each environment is described as follows:
- Hopper. In this environment, an agent needs to control a 2D single-legged robot to jump forward by controlling the torques on its joints. The action space  $A$  is  $[-1,1]^3$ , one dimension for each joint; the state space  $S$  is 11-dimensional, which describes its current angle and velocity. Note that the  $x$ -coordinate is not a part of the state, which means that the expert state approximately repeats itself periodically. Thus, in Fig. 10 the expert  $R(s)$  is periodic. For this environment as well as halfcheetah, ant, and walker2d, the reward is gained by surviving and moving forward, and the episode lasts 1,000 steps.
- Halfcheetah. Similar to hopper, an agent needs to control a 2D cheetah-like robot with torques on its joints to move forward. The action space  $A$  is  $[-1,1]^6$ , and the state space  $S$  is 17-dimensional describing its coordinate and velocity.
- Ant. In this environment, the agent controls a four-legged robotic ant to move forward in a 3D space with a 111-dimensional state space describing the coordinate and velocity of its joints, as well as contact forces on each joint. The action space is  $[-1, 1]^8$ .
- Walker2d. The agent in this environment controls a 2D two-legged robot to walk forward, with state space being 27-dimensional and action space being  $[-1,1]^8$ .
Fig. 9 shows an illustration of each environment.
Dataset Settings. We generate our dataset on the basis of SMODICE. SMODICE uses 1 trajectory (1,000 states) from the "expert-v2" dataset in D4RL [16] as the task-specific dataset  $D_{\mathrm{TS}}$  and concatenates 200 trajectories (200K state-action pairs) from the "expert-v2" dataset and the whole "random-v2" dataset in D4RL (which contains  $1M$  state-action pairs) to be the task-agnostic dataset  $D_{\mathrm{TA}}$ .
Based on this, we use the task-specific dataset from SMODICE. For the task-agnostic dataset, we take the dataset from SMODICE, concatenate all state-action pairs from all trajectories into an array, and remove a state-action pair for every  $x$  steps; we test  $x \in \{2, 3, 5, 10, 20\}$  in this work.
E.2 Offline Imitation Learning from Task-Specific Dataset with Incomplete Trajectories
Environment Settings. Sec. 4.2 tests hopper, halfcheetah, ant, and walker2d with the same environment settings as those discussed in Sec. E.1.
Dataset Settings. We use the task-agnostic dataset from SMODICE as described in Sec. E.1. For the task-specific dataset, we take the first  $x$  and last  $y$  steps from the task-specific dataset of SMODICE as the new task-specific dataset, and discard the other state-action pairs; in this work, we test  $(x,y) \in \{(1,100),(10,90),(50,50),(90,10),(100,1)\}$ .
E.3 Standard Offline Imitation Learning from Observation
Environment Settings. In addition to the four Mujoco environments tested in Sec. E.1, Sec. 4.3 tests the kitchen and antimaze environments. The detailed configuration for the two environment is described as follows:
- Kitchen. In this environment, the agent controls a 9-DoF robotic arm to complete a sequence of 4 subtasks; possible subtasks include opening the microwave, moving the kettle, turning on the light, turning on the bottom burner, turning on the top burner, opening the left cabinet, and opening the right cabinet. The state space is 60-dimensional, which includes the configuration of the robot, goal location of the items to manipulate, and current position of the items.
- Antmaze. In this environment, the agent controls a robotic ant with 29-dimensional state space and 8-dimensional action space to move from one end of a u-shaped maze to the other end.
Fig. 9 illustrates the two environments.
(a) Hopper
(c) Ant
(e) Kitchen
(b) Halfcheetah
Figure 9: Illustration of environments tested in Sec. 4.3 based on OpenAI Gym [5] and D4RL [16].
(d) Walker2d
(f) Antmaze
Dataset Settings. For the four Mujoco environments, we take the task-specific dataset from SMODICE as described in Sec. E.1. As existing methods already perform well on those environments with SMODICE's dataset settings (see Sec. F.5 for results), we introduce a more challenging task-agnostic dataset concatenating 40 trajectories (40K state-action pairs) instead of 200 from the "expert-v2" dataset and the whole "random-v2" dataset in D4RL as the task-agnostic dataset  $D_{\mathrm{TA}}$ .
For the more challenging kitchen and antimaze environments, we use identical dataset settings as SMODICE. For the kitchen environment, we use 1 trajectory (184 states) from the "kitchen-complete-v0" dataset as the task-specific dataset  $D_{\mathrm{TS}}$  and the whole "kitchen-mixed-v0" dataset as the task-agnostic dataset  $D_{\mathrm{TA}}$ , which contains 33 different task sequences with a total of 136,950 state-action pairs. For the antimaze environment, we use the expert trajectory of length 285 given by SMODICE as the task-specific dataset  $D_{\mathrm{TS}}$ . We use the dataset collected by SMODICE as the task-agnostic dataset  $D_{\mathrm{TA}}$ , which contains 1,348,687 state-action pairs.
E.4 Offline Imitation Learning from Examples
Environment Settings. Sec. 4.4 tests antimaze and kitchen in the same environment settings as the ones discussed in Sec. E.3, but we only require the agent to complete one particular subtask instead
of all four (we test "opening the microwave" and "moving the kettle"). Additionally, Sec. 4.4 tests the pointmaze environment, which is a simple environment where the agent controls a pointmass to move from the center of an empty 2D plane to a particular direction. The action is 2-dimensional; the state is 4-dimensional, which describes the current coordinate and velocity of the pointmass.
Dataset Settings. For the antmaze environment, we use the same task-agnostic dataset as the one discussed in Sec. E.3, and we use 500 expert final states selected by SMODICE as the task-specific dataset. For the kitchen environment, we use the same task-agnostic dataset as the one discussed in Sec. E.3, and randomly select 500 states from the "kitchen-mixed-v0" dataset that has the subtask of interest completed as the task-specific dataset. For the pointmaze environment, we use a trajectory of 603 expert trajectories with 50K state-action pairs as the task-agnostic dataset, where the trajectories moving up, down, left and right each take up  $1/4$ . We use the last state of all trajectories that move left as the task-specific dataset.
E.5 Offline Imitation Learning from Mismatched Dynamics
Environment Settings. Sec. 4.1 tests halfcheetah, ant, and antmaze with the same environment settings as the ones discussed in Sec. E.1 and Sec. E.3.
Dataset Settings. We use the task-agnostic dataset from Sec. E.1 (ant, halfcheetah) and Sec. E.3 (antmaze). For the task-specific dataset, we use the data generated by SMODICE, which is a single trajectory conducted by an agent with different dynamics. More specifically, for halfcheetah, we use an expert trajectory of length 1,000 by a cheetah with a much shorter torso; for ant, we use an expert trajectory of length 1,000 by an ant with one leg crippled; for antimaze, we use an expert trajectory of length 173 by a pointmass instead of an ant (and thus only the first two dimensions describing the current location are used for training the discriminator). See Appendix H of SMODICE [41] for an illustration of the environments.
F Additional Experiment Results
F.1 Can Our Method Discriminate Expert and Non-Expert Data in the Task-Agnostic Dataset?
In Sec. 4, an important problem related to our motivation is left out due to page limit: Does the algorithm really succeed in discriminating expert and non-expert trajectories and segments in the task-agnostic data? We answer the question here with Fig. 10. It shows  $R(s)$  and the weight for behavior cloning  $W(s,a)$  along an expert and a non-expert trajectory in  $D_{\mathrm{TA}}$  on halfcheetah in Sec. 4.3, hopper with first 50 step and last 50 step as task-specific dataset in Sec. 4.2, and pointmaze in Sec. 4.4. We also plot the average  $R(s)$  and weight  $W(s,a)$  for state-action pairs in all expert trajectories and non-expert trajectories in  $D_{\mathrm{TA}}$ . The result clearly shows that, even if the discriminator cannot tell the expert states from the non-expert states at the beginning of the episode because states from all trajectories are similar to each other, the weight successfully propagates from the later expert state to the early trajectory, and thus the agent knows where to go even when being close to the starting point; also, the weight  $W(s,a)$  is smoother than the raw  $R(s)$ , which prevents the agent from being lost in the middle of the trajectory because of a few steps with very small weights.
F.2 Comparison to Other Baselines
Besides the baselines tested in the main paper, we compare our method to a variety of other offline RL/IL methods. The tested methods include: a) a very recent method, ReCOIL; b) methods with extra access to expert actions or reward labels, such as DWBC [68], model-based RL methods MOReL [26] and MOPO [70]; c) offline adaption for Advantage-Weighted Regression (AWR) [45], MARWIL [62]; d) a recent Wasserstein-based offline IL method, OTR [40]. As we do not have the code for ReCOIL, and model-based RL methods takes a long time to run on our testbed, we test our method on their testbeds and use their reported numbers for comparison; for other methods, we test on our testbed (both under SMODICE setting in Sec. F.5 and standard LfO setting in Sec. 4.3). We found our method to consistently outperform all methods, even those with extra access to reward labels or expert actions.
Figure 10: Illustration of  $R(s)$ , behavior cloning weight  $W(s, a)$ , and their average on expert and non-expert trajectories in  $D_{\mathrm{TA}}$ . Curves for expert trajectories are plotted in red color, while non-expert trajectories are plotted in green color. The non-expert trajectory on the hopper environment is short because the agent topples quickly and the episode is terminated early. It is clear that our design of weights propagates the large weight from later expert states and smoothes the weight along the trajectory. Hopper has a large periodic change of  $R(s)$ , because the velocity vector for the joints of the agent is periodic.
F.2.1 ReCOIL
Similar to Sec. 4.1, we test on the four common Mujoco environments, which are hopper, halfcheetah, ant, and walker2d; the results we tested include random + expert (identical to SMODICE), random + few expert (30 expert trajectories instead of 40 in Sec. 4.3), and medium + expert (substitute the random dataset to medium dataset; see ReCOIL for details). The result of random + expert is illustrated in Fig. 26, where our method outperforms ReCOIL. The result of random + few expert is illustrated in Fig. 11, and the result of medium + expert is illustrated in Fig. 12. On all test beds, our method is significantly better than ReCOIL.
Figure 11: Performance comparison between our method and ReCOIL on the "random+few expert" testbed of ReCOIL.
Figure 12: Performance comparison between our method and ReCOIL on the "medium+expert" testbed of ReCOIL.
F.2.2 Model-Based RL
In this section, we compare our method to MOReL [26] and MOPO [70] which are provided with ground-truth reward labels, i.e., MOReL and MOPO have an advantage. The trajectory with the best return is provided as the task-specific dataset to our method. Tab. 4 shows the performance comparison between our method, MOReL and MOPO; despite being agnostic to reward labels, our method is still marginally better than MOReL (74.3 vs. 72.9 reward averaged over 9 widely tested environments {halfcheetah, hopper, walker2d}  $\times$  {medium, medium-replay, medium-expert}), and much better than MOPO (74.3 vs. 42.1 average reward).
F.2.3 DWBC, MARWIL and OTR