SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Youngsoo Jang<sup>1</sup>, Geon-Hyeong Kim<sup>1</sup>, Jongmin Lee<sup>2</sup>, Sungryull Sohn<sup>1</sup>, Byoungjip Kim<sup>1</sup>, Honglak Lee<sup>1</sup>, Moontae Lee<sup>1,3</sup>  
<sup>1</sup> LG AI Research  
<sup>2</sup> University of California, Berkeley  
<sup>3</sup> University of Illinois Chicago
We consider offline safe imitation learning (IL), where the agent aims to learn the safe policy that mimics preferred behavior while avoiding non-preferred behavior from non-preferred demonstrations and unlabeled demonstrations. This problem setting corresponds to various real-world scenarios, where satisfying safety constraints is more important than maximizing the expected return. However, it is very challenging to learn the policy to avoid constraint-violating (i.e. non-preferred) behavior, as opposed to standard imitation learning which learns the policy to mimic given demonstrations. In this paper, we present a hyperparameter-free offline safe IL algorithm, SafeDICE, that learns safe policy by leveraging the non-preferred demonstrations in the space of stationary distributions. Our algorithm directly estimates the stationary distribution corrections of the policy that imitate the demonstrations excluding the non-preferred behavior. In the experiments, we demonstrate that our algorithm learns a more safe policy that satisfies the cost constraint without degrading the reward performance, compared to baseline algorithms.
Reinforcement learning (RL) [28] aims to learn the optimal policy that maximizes the expected cumulative rewards through interacting with the environment. Recently, RL has achieved remarkable successes in various domains, but its application to real-world tasks is yet challenging mainly due to the following two reasons: First, since the interaction with the environment is costly in most real-world tasks, the policy should be learned solely from the pre-collected dataset (i.e. offline learning). Second, reward annotation, which is essential for applying the RL algorithm, is a very strong requirement that is difficult to obtain in real-world tasks.
Imitation learning (IL) [26, 1, 25] is a more realistic framework to deal with real-world tasks, where the expert's demonstrations are solely given without reward annotations. Standard IL aims to mimic the expert's demonstrations generated from the optimal policy in terms of expected return. However, in most real-world tasks, additional constraints should be considered in addition to maximizing the expected return such as the safe driving for autonomous driving and toxicity degree for the conversational agent. Fortunately, non-preferred demonstrations that violate constraints are often naturally collected. For example, black box data of an accident vehicle and toxic content reported by chatbot users are separately collected as non-preferred demonstrations. However, most existing studies have focused only on imitating expert (i.e. preferred) behavior, and a safe imitation learning method that avoids non-preferred behavior has not yet been explored.
In this paper, we are particularly interested in solving safe imitation learning in an offline setting: Given scarce but labeled non-preferred demonstrations, and abundant but unlabeled demonstrations
Figure 1: Illustrative example of our problem setting and the position of work among the prior methods and problem settings. First, prior offline imitation learning algorithms (DemoDICE [13], DWBC [31]), which aim to imitate expert demonstrations, assume labeled preferred demonstrations and unlabeled demonstrations. Second, preference-based policy learning (PPL) algorithms and DExperts [19] require both labels of preferred and non-preferred demonstrations. Unlike the problem settings of prior works, our work focuses on the problem of safe imitation learning assuming only the labels of non-preferred demonstration that can deal with safety more importantly.
of both preferred and non-preferred behavior, we aim to learn a safe policy that follows the preferred behavior while avoiding the non-preferred behavior (see Figure 1). This problem setting has been mostly unexplored, and cannot be completely addressed by existing methods. One of the existing applicable approaches is the preference-based policy learning algorithm [4, 35, 17], which optimizes a policy based on preference rankings over demonstrations. This method can learn the policy to prefer unlabeled demonstrations to non-preferred demonstrations, but it also prefers non-preferred demonstrations mixed in the unlabeled demonstrations, which can lead to learning suboptimal policy. One of the approaches effectively using the unlabeled dataset is DWBC [31], a weighted behavior cloning based on a discriminator, which adopts the positive unlabeled learning [9, 36, 30] for training the discriminator. This method requires an additional cost of hyperparameter search for training the discriminator, and its performance is hyperparameter sensitive. Another notable approach in the natural language processing domains is DExperts [19], a decoding-time method, which combines trained models with preferred and non-preferred demonstrations. This method is also hyperparameter sensitive and prone to the degeneration issue that leads to abnormal behavior. Both DWBC and DExperts, which require hyperparameter search, are unsuitable for real-world tasks where interaction with the real environment is costly and risky.
We present an offline safe IL algorithm that learns safe policy by leveraging the non-preferred demonstrations in the space of stationary distributions. Our algorithm, offline Safe imitation learning by avoiding non-preferred demonstrations via DDistribution Correction Estimation (SafeDICE), directly estimates the stationary distribution corrections of the policy that imitate the demonstrations excluding the non-preferred behavior. SafeDICE essentially leverages the non-preferred demonstrations in the space of stationary distributions, unlike existing algorithms that leverage it in the space of policy [19] or the learning process of the discriminator [31]. We show that our algorithm can be reduced to a single convex minimization problem without additional hyperparameter search, unlike existing algorithms that are highly sensitive to the hyperparameter. In the experiments, we demonstrate that our algorithm learns a more preferred (i.e. constraint-satisfying) policy compared to baseline algorithms in constrained RL benchmarks including RWRL [8] and Safety Gym [24].
2.1 Markov Decision Process (MDP)
We consider an environment modeled as a Markov Decision Process (MDP), defined by  $M = \langle S, A, T, R, p_0, \gamma \rangle$  [28], where  $S$  is the set of states,  $A$  is the set of actions,  $T: S \times A \to \Delta(S)$  is the transition probability,  $R: S \times A \to \mathbb{R}$  is the reward function,  $p_0 \in \Delta(S)$  is the distribution of the initial state, and  $\gamma \in [0, 1)$  is the discount factor. For the given policy  $\pi$ , its stationary distribution  $d^{\pi}(s, a)$  is defined as follows:
$$
d ^ {\pi} (s, a) := (1 - \gamma) \sum_ {t = 0} ^ {\infty} \gamma^ {t} \Pr (s _ {t} = s, a _ {t} = a)
$$
where  $s_0 \sim p_0, a_t \sim \pi(\cdot|s_t)$ , and  $s_{t+1} \sim T(\cdot|s_t, a_t)$  for all timesteps  $t \geq 0$ . For brevity,  $d^\pi$  will be used to denote the stationary distribution  $d^\pi(s, a)$ .
We consider the offline imitation learning setting, where the policy should be optimized from solely a pre-collected dataset without reward annotation and any additional online environment interaction. In our problem setting, preferred and non-preferred behaviors are defined by satisfying or violating constraints. We denote the preferred policy by  $\pi^P$ , the non-preferred policy by  $\pi^N$ , and the corresponding stationary distribution of  $\pi^P$  and  $\pi^N$  by  $d^P$  and  $d^N$ , respectively. We assume a pre-collected dataset that consists of scarce but labeled dataset  $D^N$  generated by the non-preferred policy, and a abundant but unlabeled dataset  $D^U$  generated by both preferred and non-preferred policies. We denote the corresponding stationary distribution of  $D^U$  by  $d^U$ , which can be regarded as convex combination of  $d^P$  and  $d^N$  with unknown ratio  $\alpha$  (i.e.  $d^U \coloneqq (1 - \alpha)d^P + \alpha d^N$ ). In this setting, we aim to learn a safe policy from unlabeled demonstrations while avoiding non-preferred demonstrations that are labeled as non-preferred behavior.
2.2 Preference-based Offline Imitation Learning
Imitation learning (IL) aims to imitate the behavior of expert policy solely from state-action demonstrations without reward annotations. The standard IL is naturally formulated as minimizing the KL divergence between stationary distributions of preferred policy (i.e. expert policy) and target policy [14, 11]:
$$
\pi^ {*} := \underset {\pi} {\arg \min } D _ {\mathrm {K L}} \left(d ^ {\pi} \| d ^ {P}\right). \tag {1}
$$
However, it requires a large amount of annotated preferred demonstrations, which may be a too strong requirement for various real-world scenarios. Recent offline IL method [13] relaxes this requirement for preferred demonstrations by leveraging the unlabeled dataset which may contain both preferred and non-preferred demonstrations:
$$
\pi^ {*} := \underset {\pi} {\arg \min } D _ {\mathrm {K L}} \left(d ^ {\pi} \| d ^ {P}\right) + \alpha D _ {\mathrm {K L}} \left(d ^ {\pi} \| d ^ {U}\right), \tag {2}
$$
where  $\alpha \geq 0$  is a hyperparameter that controls the balance between minimizing the KL divergence with  $d^{U}$  and  $d^{P}$ . This offline imitation learning method can learn the preferred policy from a large number of unlabeled demonstrations with only a small amount of labeled preferred demonstrations.
In contrast to the mainstream of recent IL approaches focusing on mimicking expert (i.e. preferred) behavior, avoiding dangerous (i.e. non-preferred) behavior is a more important problem in many real-world tasks, which has been unexplored yet. In the same manner as Eq. (2), we can naively consider the penalty for the non-preferred behavior by maximizing the KL divergence with  $d^{N}$ :
$$
\pi^ {*} := \underset {\pi} {\arg \min } - D _ {\mathrm {K L}} \left(d ^ {\pi} \| d ^ {N}\right) + \alpha D _ {\mathrm {K L}} \left(d ^ {\pi} \| d ^ {U}\right), \tag {3}
$$
where  $\alpha \geq 0$  is a hyperparameter that controls the balance between minimizing the KL divergence with  $d^{U}$  and maximizing the KL divergence with  $d^{N}$ . However, since the KL divergence has no upper bound, optimizing Eq. (3) is very unstable and challenging, where the  $D_{\mathrm{KL}}(d^{\pi}||d^{N})$  can blow up to infinity. Therefore, safe imitation learning while avoiding non-preferred behavior is not a straightforward problem and has not yet been studied.
3 SafeDICE
In this section, we present offline Safe imitation learning by avoiding the non-preferred demonstrations via Distribution Correction Estimation (SafeDICE), an offline IL algorithm that can learn the safe policy using labeled non-preferred demonstrations. SafeDICE optimizes the stationary distributions of the target policy to match the stationary distribution of the preferred policy from the unlabeled dataset by excluding non-preferred demonstrations.
3.1 Offline Safe Imitation Learning with Non-Preferred Demonstrations
Similar to standard IL approaches, the derivation of our algorithm starts from KL divergence minimization between  $d^{\pi}$  and  $d^{P}$ , which can be estimated by  $d^{U}$  and  $d^{N}$  as follows:
$$
\pi^ {*} = \underset {\pi} {\arg \min } D _ {\mathrm {K L}} \left(d ^ {\pi} \left\| \frac {d ^ {U} - \alpha d ^ {N}}{1 - \alpha}\right), \right. \tag {4}
$$
where  $\alpha$  is unknown ratio between the stationary distributions of preferred policy and non-preferred policy (i.e.  $d^{U} \coloneqq (1 - \alpha)d^{P} + \alpha d^{N}$ ). Optimizing the Eq. (4) is sensitive to  $\alpha$ , and selecting the  $\alpha$  with hyperparameter search is not suitable for real-world tasks where the interaction with the real environment is very costly and risky. Therefore, it is not straightforward to learn the preferred policy using non-preferred demonstrations with existing imitation learning algorithms. The main contribution of our paper is to present an algorithm that theoretically determines  $\alpha$  and learns the preferred policy only with the labels of non-preferred demonstrations without any hyperparameter search. In the following description, we first consider the  $\alpha$  as a known ratio in our derivation, then present the way to select the  $\alpha$  while providing the theoretical guarantees on the error bound of optimized stationary distribution with selected  $\alpha$ . All the proofs can be found in Appendix A.
The main flow of our derivation follows prior DICE-based offline imitation learning methods [12, 13], but it is clearly different from the existing methods that require hyperparameter search to deal with unknown  $\alpha$ . First, we consider a problem equivalent to Eq. (4) in terms of stationary distribution  $d$ :
$$
\max  _ {d \geq 0} - D _ {\mathrm {K L}} \left(d \left\| \frac {d ^ {U} - \alpha d ^ {N}}{1 - \alpha}\right) \right. \tag {5}
$$
$$
\text {s . t .} (\mathcal {B} _ {*} d) (s) = (1 - \gamma) p _ {0} (s) + \gamma (\mathcal {T} _ {*} d) (s) \quad \forall s, \tag {6}
$$
where  $(\mathcal{B}_*d)(s) := \sum_a d(s,a)$  is the marginalization operator, and  $(\mathcal{T}_*d)(s) := \sum_{\bar{s},\bar{a}} T(s|\bar{s},\bar{a}) d(\bar{s},\bar{a})$  is the transposed Bellman operator. The Bellman flow constraint (6) ensures that  $d(s,a)$  is a valid stationary distribution of some policy, where  $d(s,a)$  can be interpreted as a normalized occupancy measure of  $(s,a)$ . Therefore, if we find optimal  $d^*(s,a)$  that satisfies Eq. (5) and (6), then its corresponding policy can be easily obtained by  $\pi^*(a|s) = \frac{d^*(s,a)}{\sum_a d^*(s,a)}$ .
The Lagrangian dual formulation of the above constrained optimization problem is
$$
\max  _ {d \geq 0} \min  _ {\nu} - D _ {\mathrm {K L}} \left(d \left\| \frac {d ^ {U} - \alpha d ^ {N}}{1 - \alpha}\right) + \sum_ {s} \nu (s) \left((1 - \gamma) p _ {0} (s) + \gamma (\mathcal {T} _ {*} d) (s) - (\mathcal {B} _ {*} d) (s)\right), \quad (7) \right.
$$
where  $\nu(s)$  are the Lagrange multipliers for the Bellman flow constraints (6). However, directly optimizing the Eq. (7) is challenging in an offline manner, where further environment interactions are not available. To obtain a tractable objective, we introduce the following derivations:
$$
\begin{array}{l} - D _ {\mathrm {K L}} \left(d \left\| \frac {d ^ {U} - \alpha d ^ {N}}{1 - \alpha}\right) + \sum_ {s} \nu (s) \left((1 - \gamma) p _ {0} (s) + \gamma (\mathcal {T} _ {*} d) (s) - (\mathcal {B} _ {*} d) (s)\right) \right. \\ = (1 - \gamma) \mathbb {E} _ {p _ {0}} [ \nu (s) ] + \mathbb {E} _ {d} \bigg [ \gamma (\mathcal {T} \nu) (s, a) - \nu (s) - \log \frac {(1 - \alpha) d (s , a)}{d ^ {U} (s , a) - \alpha d ^ {N} (s , a)} \bigg ] \\ = (1 - \gamma) \mathbb {E} _ {p _ {0}} [ \nu (s) ] + \mathbb {E} _ {d} \left[ \gamma (\mathcal {T} \nu) (s, a) - \nu (s) - \log \underbrace {\frac {d (s , a)}{d ^ {U} (s , a)}} _ {:= w (s, a)} + \underbrace {\log \frac {d ^ {U} (s , a) - \alpha d ^ {N} (s , a)}{(1 - \alpha) d ^ {U} (s , a)}} _ {:= r (s, a)} \right] \\ = (1 - \gamma) \mathbb {E} _ {p _ {0}} [ \nu (s) ] + \mathbb {E} _ {d} \left[ \underbrace {r _ {\alpha} (s , a) + \gamma (\mathcal {T} \nu) (s , a) - \nu (s)} _ {:= A _ {\nu} (s, a)} - \log w (s, a) \right] \\ = (1 - \gamma) \mathbb {E} _ {p _ {0}} [ \nu (s) ] + \mathbb {E} _ {d ^ {U}} \left[ w (s, a) \left(A _ {\nu} (s, a) - \log w (s, a)\right) \right] \\ =: \mathcal {L} (w, \nu ; r _ {\alpha}), \\ \end{array}
$$
where  $r_{\alpha}(s,a)$  corresponds to the preference-based reward function in the context of RL, which will be trained by the pre-collected dataset.
In summary, our goal is reduced to solve the following maximin optimization:
$$
\max  _ {w \geq 0} \min  _ {\nu} \mathcal {L} (w, \nu ; r _ {\alpha}), \tag {8}
$$
where the optimal solution  $w^{*}$  of Eq. (8) represents the stationary distribution correction of an optimal policy  $\pi^{*}$  (i.e.  $w^{*}(s,a) = \frac{d^{\pi^{*}}(s,a)}{d^{U}(s,a)}$ ).
3.2 Pretraining the Preference-based Reward Function  $r_{\alpha}(s,a)$
To optimize the Eq. (8), we pretrain the discriminator to estimate the log ratio of  $d^{N}(s,a)$  and  $d^{U}(s,a)$ , which is required to estimate for the  $r_{\alpha}(s,a)$ :
$$
c ^ {*} = \underset {c} {\arg \max } \mathbb {E} _ {d ^ {N}} [ \log c (s, a) ] + \mathbb {E} _ {d ^ {U}} [ \log (1 - c (s, a)) ]. \tag {9}
$$
With the optimal discriminator  $c^*(s, a) = \frac{d^N(s, a)}{d^U(s, a) + d^N(s, a)}$  trained by Eq. (9), the preference-based reward function  $r_\alpha(s, a)$  can also be obtained as:
$$
r _ {\alpha} (s, a) = \log \left(\frac {1 - (1 + \alpha) c ^ {*} (s , a)}{(1 - \alpha) (1 - c ^ {*} (s , a))}\right). \tag {10}
$$
However, according to the value of  $\alpha$ , the reward function  $r_{\alpha}(s,a)$  defined by Eq. (10) cannot guarantee that the value inside the log function always becomes a positive. Therefore, finding the optimal  $\alpha$  through hyperparameter search is infeasible.
3.3 Reduction to the Single Min Optimization
Since the optimization Eq. (7) is a convex optimization, we can change the order of operator from maximin optimization to minimax optimization [2]:
$$
\min  _ {\nu} \max  _ {w \geq 0} \mathcal {L} (w, \nu ; r). \tag {11}
$$
Then, we can replace the inner maximization of Eq. (11) with a closed-form solution for  $w$ , which can be easily obtained by exploiting the convexity.
Proposition 3.1. For any  $\nu$ , the closed-form solution to the inner maximization of Eq. (11), i.e.  $w_{\nu}^{*} = \arg \max_{w\geq 0}\mathcal{L}(w,\nu ;r)$ , is given by:
$$
w _ {\nu} ^ {*} (s, a) = \exp \left(A _ {\nu} (s, a) - 1\right). \tag {12}
$$
By using the closed-form solution  $w_{\nu}^{*}$  of Eq. (12), our derivation can be reduced to a single minimization as follows:
$$
\min  _ {\nu} \mathcal {L} \left(w _ {\nu} ^ {*}, \nu ; r\right) = (1 - \gamma) \mathbb {E} _ {p _ {0}} [ \nu (s) ] + \mathbb {E} _ {d U} \left[ \exp \left(A _ {\nu} (s, a) - 1\right) \right]. \tag {13}
$$
Finally, by optimizing the single minimization problem of Eq. (13), we can obtain the optimal stationary distribution correction  $w^{*}$  which can easily extract the corresponding optimal policy  $\pi^{*}$ . To extract a policy from the optimal stationary distribution correction  $w^{*}(s,a)$ , we adopt weighted behavior cloning on the offline dataset as follows:
$$
\min  _ {\pi} - \mathbb {E} _ {(s, a) \sim d ^ {\pi^ {*}}} [ \log \pi (a | s) ] = - \mathbb {E} _ {(s, a) \sim d ^ {U}} \left[ w ^ {*} (s, a) \log \pi (a | s) \right], \tag {14}
$$
where  $w^{*}(s,a) = \frac{d^{\pi^{*}}(s,a)}{d^{U}(s,a)}$
3.4 Hyperparameter  $\alpha$  Selection
In the previous sections, we have derived our algorithm under the assumption that  $\alpha$  is a known ratio. The remaining question is how we select the suitable  $\alpha$  without additional hyperparameter search. Fortunately, in contrast to the existing methods [19, 31] that require the additional cost of hyperparameter search, we can specify the  $\alpha$  with providing the theoretical guarantees on the error bound of the estimated stationary distribution.
Proposition 3.2. Suppose  $\min_{(s,a)}\frac{d^P(s,a)}{d^N(s,a)}\leq \epsilon$ . If we obtain the estimated stationary distribution by optimizing the Eq. (7) with  $\alpha$  set to  $\hat{\alpha}$  as follows:
$$
\hat {\alpha} = \min  _ {(s, a)} \frac {d ^ {U} (s , a)}{d ^ {N} (s , a)}, \tag {15}
$$
then the total variation distance of true and estimated stationary distribution of preferred policy is bounded as follows:
$$
D _ {\mathrm {T V}} \left(d ^ {P} (s, a), \hat {d} ^ {P} (s, a)\right) \leq \frac {\epsilon}{1 + \epsilon}, \tag {16}
$$
where  $D_{\mathrm{TV}}$  is total variation distance,  $d^{P}(s,a)$  is true stationary distribution of the preferred policy and  $\hat{d}^P (s,a)$  is estimated stationary distribution which is defined by  $\hat{\alpha}$  as  $\hat{d}^P (s,a) = \frac{d^U(s,a) - \hat{\alpha}d^N(s,a)}{1 - \hat{\alpha}}$ . (Proof in Appendix A.)
Intuitively, since the behavior of the non-preferred and preferred policies are different, there must exist state-action spaces that are rarely visited by the preferred policy but are frequently visited by the non-preferred policy. Therefore  $\epsilon$  may be a very small positive value, then the total variation distance of the true and estimated stationary distribution is close to zero:
$$
\lim  _ {\epsilon \rightarrow 0} D _ {\mathrm {T V}} \left(d ^ {P} (s, a), \hat {d} ^ {P} (s, a)\right) = 0. \tag {17}
$$
We also provide proof in Appendix A.3 that  $r_{\alpha}(s,a)$  of Eq. (10), defined by the  $\alpha$  selected by Eq. (15), is always well-defined with positive value inside the log-function in Eq. (10).
Finally, our algorithm optimizes the Eq. (7) with the  $\alpha$  obtained by Eq. (15), where  $\alpha$  is no longer an unknown hyperparameter. In short, by leveraging the non-preferred demonstrations in the space of stationary distributions, this problem can be resolved by solving a single convex minimization problem for a specific  $\alpha$ . This is in contrast to the existing methods, which directly leverage the non-preferred demonstrations for training the policy [19] or discriminator [31] that require the additional cost of hyperparameter search. The pseudocode for the whole process of SafeDICE can be found in Appendix B.
Stationary distribution correction estimation (DICE) Since the introduction of DICE for off-policy evaluation in DualDICE [21], various DICE-family algorithms have been developed for off-policy evaluation [33, 34, 32, 7]. In addition to these studies, two algorithms, AlgaeDICE [22] and ValueDICE [14], have been proposed to address the challenges of off-policy RL and IL, respectively. However, they may encounter issues with numerical instability in an offline setting due to the nested min-max optimization and out-of-distribution action evaluation.
To overcome the issues of numerical instability in offline settings, recent DICE-based offline RL [15, 16] and IL algorithms [13, 20, 12] have adopted a strategy of directly optimizing stationary distribution correction ratios, similar to the approach used in our work. Additionally, these offline IL algorithms utilize unlabeled demonstrations in an effort to address the problem of distribution drift [25]. They share a similarity with our algorithm in that they all aim to recover the expert's behavior through direct optimization of the stationary distribution correction ratios. However, these algorithms rely on the availability of clean expert demonstrations, which can make them difficult to use in situations where expert demonstrations are not available.
Policy Learning from Suboptimal Demonstrations Preference-based policy learning algorithms use ranked suboptimal demonstrations to learn a reward function. Many of these algorithms, such as T-REX [4], PrefPPO [18], and PEBBLE [17], utilize Bradley-Terry model [3] to model the reward function. Once the reward function is learned, T-REX and PrefPPO use PPO [27] to train a policy, while PEBBLE employs SAC [10]. While preference-based policy learning algorithms allow for learning from unlabeled demonstrations and potentially disregarding the negative influence of non-preferred demonstrations, it is important to note that they may also give preference to non-preferred demonstrations mixed in the unlabeled demonstrations.
On the other hand, there are also many studies that utilize suboptimal demonstrations, even though they are not preference-based policy learning algorithms. D-REX [5] and SSRR [6] both learn a policy by using given suboptimal demonstrations, and then generate trajectories with varying levels of optimality using this policy with injected noise. But, both algorithms require online sampling, which is not feasible in offline settings. Wu et al. [29] proposed two algorithms that utilize positive-unlabeled learning [9, 36, 30] to train a discriminator for constructing a reward function. However,