SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Table 9: CEIL hyper-parameters.
Parameter	Value
size of the embedding dictionary	4096
size of the embedding dimension	16
trajectory window size	2
encoder: optimizer	Adam
encoder: learning rate	3e-4
encoder: learning rate scheduler	CosineAnnealingWarmRestarts(T_0 = 1000,TMULT=1, eta_min=1e-5)
encoder: number of hidden layers	4
encoder: number of hidden units per layer	512
encoder: nonlinearity	ReLU
policy: optimizer	Adam
policy: learning rate	3e-4
policy: learning rate scheduler	CosineAnnealingWarmRestarts(T_0 = 1000,TMULT=1, eta_min=1e-5)
policy: number of hidden layers	4
policy: number of hidden units per layer	512
policy: nonlinearity	ReLU
decoder: optimizer	Adam
decoder: learning rate	3e-4
decoder: learning rate scheduler	CosineAnnealingWarmRestarts(T_0 = 1000,TMULT=1, eta_min=1e-5)
decoder: number of hidden layers	4
decoder: number of hidden units per layer	512
decoder: nonlinearity	ReLU
Table 10: Baseline methods and their code-bases.
Baselines	Code-bases
GAIL, GAIfo, AIRL	https://github.com/HumanCompatibleAI/imitation
SAIL	https://github.com/FangchenLiu/SAIL
IQ-Learn, SQL	https://github.com/Div99/IQ-Learn
ValueDICE	https://github.com/google-research/google-research/tree/master/value_dice
DemoDICE	https://github.com/KAIST-AILab/imitation-dice
SMODICE, ORIL	https://github.com/JasonMa2016/SMODICE
9.4 Baselines Implementation Details
We summarize our code-bases of our baseline implementations in Table 10 and describe each baseline as follows:
Generative Adversarial Imitation Learning (GAIL). GAIL [30] is a GAN-based online LfD method that trains a policy (generator) to confuse a discriminator trained to distinguish between generated transitions and expert transitions. While the goal of the discriminator is to maximize the
objective below, the policy is optimized via an RL algorithm to match the expert occupancy measure (minimize the objective below):
$$
\mathcal {J} (\pi , D) = \mathbb {E} _ {\pi} [ \log (D (s, a)) ] + \mathbb {E} _ {\pi_ {E}} [ 1 - \log (D (s, a)) ] - \lambda H (\pi).
$$
We used the implementation by Gleave et al. [29] on the GitHub page<sup>8</sup>, where there are two modifications introduced with respect to the original paper: 1) a higher output of the discriminator represents better, 2) PPO is used to optimize the policy instead of TRPO.
Generative Adversarial Imitation from Observations (GAIfO). GAIfO [66] is an online LfO method that applies the principle of GAIL and utilizes a state-only discriminator to judge whether the generated trajectory matches the expert trajectory in terms of states. We provide the objective of GAIfO as follows:
$$
\mathcal {J} (\pi , D) = \mathbb {E} _ {\pi} [ \log (D (s, s ^ {\prime})) ] + \mathbb {E} _ {\pi_ {E}} \left[ 1 - \log (D (s, s ^ {\prime})) \right] - \lambda H (\pi).
$$
Based on the implementation of GAIL, we implement GAIfO by changing the input of the discriminator to state transitions.
Adversarial Inverse Reinforcement Learning (AIRL). AIRL [21] is an online LfD/LfO method using an adversarial learning framework similar to GAIL. It modifies the form of the discriminator to explicitly disentangle the task-relevant information from the transition dynamics. To make the policy more generalized and less sensitive to dynamics, AIRL proposes to learn a parameterized reward function using the output of the discriminator:
$$
f _ {\theta , \phi} (s, a, s ^ {\prime}) = g _ {\theta} (s, a) + \lambda h _ {\phi} (s ^ {\prime}) - h _ {\phi} (s),
$$
$$
D _ {\theta , \phi} (s, a, s ^ {\prime}) = \frac {\exp (f _ {\theta , \phi} (s , a , s ^ {\prime}))}{\exp (f _ {\theta , \phi} (s , a , s ^ {\prime})) + \pi (a | s)}.
$$
Similarly to GAIL, we used the code provided by Gleave et al. [29], and the RL algorithm is also PPO.
State Alignment-based Imitation Learning (SAIL). SAIL [48] is an online LfO method capable of solving cross-domain tasks. SAIL aims to minimize the divergence between the policy rollout and the expert trajectory from both local and global perspectives: 1) locally, a KL divergence between the policy action and the action predicted by a state planner and an inverse dynamics model, 2) globally, a Wasserstein divergence of state occupancy between the policy and the expert. The policy is optimized using:
$$
\begin{array}{l} \mathcal {J} (\pi) = - D _ {\mathcal {W}} (\pi (s) \| \pi_ {E} (s)) - \lambda D _ {K L} (\pi (\cdot | s _ {t}) \| \pi_ {E} (\cdot | s _ {t})) \\ = \mathbb {E} _ {\pi (s _ {t}, a _ {t}, s _ {t + 1})} \bigg (\sum_ {t = 1} ^ {T} \frac {D (s _ {t + 1}) - \mathbb {E} _ {\pi_ {E} (s)} D (s)}{T} \bigg) - \lambda D _ {K L} \bigg (\pi \big (\cdot | s _ {t} \big) \| g _ {\mathrm {i n v}} \big (\cdot | s _ {t}, f (s _ {t}) \big) \bigg), \\ \end{array}
$$
where  $D$  is a state-based discriminator trained via  $\mathcal{J}(D) = \mathbb{E}_{\pi_E}[D(s)] - \mathbb{E}_\pi [D(s)]$ ,  $f$  is the pretrained VAE-based state planner, and  $g_{\mathrm{inv}}$  is the inverse dynamics model trained by supervised regression.
In the online setting, we use the official implementation published by the authors<sup>9</sup>, where SAIL is optimized using PPO with the reward definition:  $r(s_{t}, s_{t+1}) = \frac{1}{T} \left[ D(s_{t+1}) - \mathbb{E}_{\pi_{E}(s)} D(s) \right]$ . Besides, we further implement SAIL in the offline setting by using TD3+BC [23] to maximize the reward defined above.
In our experiments, we empirically discover that SAIL is computationally expensive. While SAIL is able to learn tasks in the typical IL setting ( $S$ -on- $LfD$ ), our early experimental results find that SAIL(TD3+BC) with heavy hyperparameter tuning failed on the offline setting. This indicates that SAIL is rather sensitive to the dataset composition, which also coincides with the results gathered in Ma et al. [56]. Thus, we do not include SAIL in our comparison results.
Soft-Q Imitation Learning (SQIL). SQIL [61] is a simple but effective single-domain LfD IL algorithm that is easy to implement with both online and offline Q-learning algorithms. The main
$^{8}$ https://github.com/HumanCompatibleAI/imitation
<sup>9</sup>https://github.com/FangchenLiu/SAIL
idea of SQIL is to give sparse rewards  $(+1)$  only to those expert transitions and zero rewards (0) to those experiences in the replay buffer. The Q-function of SQIL is updated using the squared soft Bellman Error:
$$
\delta^ {2} (\mathcal {D}, r) \triangleq \frac {1}{| \mathcal {D} |} \sum_ {(s, a, s ^ {\prime}) \in \mathcal {D}} \left(Q (s, a) - \left(r + \gamma \log \big (\sum_ {a ^ {\prime} \in \mathcal {A}} \exp (Q (s ^ {\prime}, a ^ {\prime})) \big)\right)\right) ^ {2}.
$$
The overall objective of the Q-function is to maximize the following objective:
$$
\mathcal {J} (Q) = - \delta^ {2} (\mathcal {D} _ {E}, 1) - \delta^ {2} (\mathcal {D} _ {\pi}, 0).
$$
In our experiments, the online imitation policy is optimized using SAC which is also used in the original paper. To make a fair comparison among the offline IL baselines, the offline policy is optimized via TD3+BC.
Offline Reinforced Imitation Learning (ORIL). ORIL [78] is an offline single-domain IL method that solves both LfD and LfO tasks. To relax the hard-label assumption (like the sparse rewards made in SQIL), ORIL treats the experiences stored in the replay buffer as unlabelled data that could potentially include both successful and failed trajectories. More specifically, ORIL aims to train a reward function to distinguish between the expert and the suboptimal data without explicitly knowing the negative labels. By incorporating Positive-unlabeled learning (PU-learning), the objective of the reward model can be written as follows (for the LfD setting):
$$
\mathcal {J} (R) = \eta \mathbb {E} _ {\pi_ {E} (s, a)} [ \log (R (s, a)) ] + \mathbb {E} _ {\pi (s, a)} [ \log (1 - R (s, a)) ] - \eta \mathbb {E} _ {\pi_ {E} (s, a)} [ \log (1 - R (s, a)) ],
$$
where  $\eta$  is the relative proportion of the expert data and we set it as 0.5 throughout our experiments. In the original paper, the policy learning algorithm of ORIL is Critic Regularized Regression (CRR), while in this paper, we implemented ORIL using TD3+BC for fair comparisons. Besides, we adapted ORIL to the LfO setting by learning a state-only reward function:
$$
\mathcal {J} (R) = \eta \mathbb {E} _ {\pi_ {E} (s, s ^ {\prime})} [ \log (R (s, s ^ {\prime})) ] + \mathbb {E} _ {\pi (s, s ^ {\prime})} [ \log (1 - R (s, s ^ {\prime})) ] - \eta \mathbb {E} _ {\pi_ {E} (s, s ^ {\prime})} [ \log (1 - R (s, s ^ {\prime})) ].
$$
Inverse soft-Q learning (IQ-Learn). IQ-Learn [28] is an IRL-based method that can solve IL tasks in the online/offline and LfD/LfO settings. It proposes to directly learn a Q-function from demonstrations and avoid the intermediate step of reward learning. Unlike GAIL optimizing a min-max objective defined in the reward-policy space, IQ-Learn solves the expert matching problem directly in the policy-Q space. The Q-function is trained to maximize the objective:
$$
\mathbb {E} _ {\pi_ {E} (s, a, s ^ {\prime})} \left[ Q (s, a) - \gamma V ^ {\pi} (s ^ {\prime}) \right] - \mathbb {E} _ {\pi (s, a, s ^ {\prime})} \left[ Q (s, a) - \gamma V ^ {\pi} (s ^ {\prime}) \right] - \psi (r),
$$
where  $V^{\pi}(s) \triangleq \mathbb{E}_{a \sim \pi(\cdot | s)}[Q(s, a) - \log \pi(a | s)]$ ,  $\psi(r)$  is a regularization term calculated over the expert distribution. Then, the policy is learned by SAC.
We use the code provided in the official IQ-learn repository $^{10}$  and reproduce the online-LfD results reported in the original paper. For online tasks, we empirically find that penalizing the Q-value on the initial states gives the best and most stabilized performance. The learning objective of the Q-function for the online tasks is:
$$
\mathcal {J} (Q) = \mathbb {E} _ {\pi_ {E} (s, a, s ^ {\prime})} [ Q (s, a) - \gamma V ^ {\pi} (s ^ {\prime}) ] - (1 - \gamma) \mathbb {E} _ {\rho_ {0}} [ V ^ {\pi} (s _ {0}) ] - \psi (r).
$$
In the offline setting, we find that using the above objective easily leads to an overfitting issue, causing collapsed performance. Thus, we follow the instruction provided in the paper and only penalize the expert samples:
$$
\begin{array}{l} \mathcal {J} (Q) = \mathbb {E} _ {\pi_ {E} (s, a, s ^ {\prime})} [ Q (s, a) - \gamma V ^ {\pi} (s ^ {\prime}) ] - \mathbb {E} _ {\pi_ {E} (s, a, s ^ {\prime})} \left[ V ^ {\pi} (s) - \gamma V ^ {\pi} (s ^ {\prime}) \right] - \psi (r) \\ = \mathbb {E} _ {\pi_ {E} (s, a, s ^ {\prime})} [ Q (s, a) - V ^ {\pi} (s) ] - \psi (r). \\ \end{array}
$$
Imitation Learning via Off-Policy Distribution Matching (ValueDICE). ValueDICE [41] is a DICE-based<sup>11</sup> LfD algorithm which minimizes the divergence of state-action distributions between the policy and the expert. In contrast to the state-conditional distribution of actions  $\pi(\cdot|s)$  used in the
<sup>10</sup>https://github.com/Div99/IQ-Learn
$^{11}$ DICE refers to stationary Distribution Estimation Correction
above methods, the state-action distribution,  $d^{\pi}(s,a):\mathcal{S}\times \mathcal{A}\to [0,1]$ , can uniquely characterize a one-to-one correspondence,
$$
d ^ {\pi} (s, a) \triangleq (1 - \gamma) \sum_ {t = 0} ^ {\infty} \gamma^ {t} \mathbf {P r} (s _ {t} = s, a _ {t} = a | s _ {0} \sim \rho_ {0}, a _ {t} \sim \pi (s _ {t}), s _ {t + 1} \sim P (s _ {t}, a _ {t})).
$$
Thus, the plain expert matching objective can be reformulated and expressed in the Donsker-Varadhan representation:
$$
\begin{array}{l} \mathcal {J} (\pi) = - D _ {K L} \left(d ^ {\pi} (s, a) \| d ^ {\pi_ {E}} (s, a)\right) \\ = \min  _ {x: \mathcal {S} \times \mathcal {A} \to \mathbb {R}} \log \mathbb {E} _ {(s, a) \sim d ^ {\pi} E} [ \exp (x (s, a)) ] - \mathbb {E} _ {(s, a) \sim d ^ {\pi}} [ x (s, a) ]. \\ \end{array}
$$
The objective above can be expanded further by defining  $x(s,a) = v(s,a) - \mathcal{B}^{\pi}v(s,a)$  and using a zero-reward Bellman operator  $\mathcal{B}^{\pi}$  to derive the following (adversarial) objective:
$$
\mathcal {J} _ {D I C E} (\pi , v) = \log \mathbb {E} _ {(s, a) \sim d ^ {\pi_ {E}}} \left[ \exp \left(v (s, a) - \mathcal {B} ^ {\pi} v (s, a)\right) \right] - (1 - \gamma) \mathbb {E} _ {s _ {0} \sim \rho_ {0}, a _ {0} \sim \pi (\cdot | s _ {0})} \left[ v (s _ {0}, a _ {0}) \right].
$$
We use the official Tensorflow implementation<sup>12</sup> in our experiments. In the online setting, the rollouts collected are used as an additional replay regularization. The overall objective in the online setting is:
$$
\begin{array}{l} \mathcal {J} _ {D I C E} ^ {m i x} (\pi , v) \\ = - D _ {K L} \left((1 - \alpha) d ^ {\pi} (s, a) + \alpha d ^ {R B} (s, a) \| (1 - \alpha) d ^ {\pi_ {E}} (s, a) + \alpha d ^ {R B} (s, a)\right) \\ = \log \mathbb {E} _ {(s, a) \sim d ^ {m i x}} [ \exp (v (s, a) - \mathcal {B} ^ {\pi} v (s, a)) ] - (1 - \alpha) (1 - \gamma) \mathbb {E} _ {s _ {0} \sim \rho_ {0}, a _ {0} \sim \pi (\cdot | s _ {0})} [ v (s _ {0}, a _ {0}) ] \\ - \alpha \mathbb {E} _ {(s, a) \sim d ^ {R B}} [ v (s, a) - \mathcal {B} ^ {\pi} v (s, a) ], \\ \end{array}
$$
where  $d^{mix} \triangleq (1 - \alpha)d^{\pi_E} + \alpha d^{RB}$  and  $\alpha$  is a non-negative regularization coefficient (we set  $\alpha$  as 0.1 following the specification of the paper).
In the offline setting, ValueDICE only differs in the source of sampling data. We change the online replay buffer to the offline pre-collected dataset.
Offline Imitation Learning with Supplementary Imperfect Demonstrations (DemoDICE).
DemoDICE [39] is a DICE-based offline LfD method that assumes to have access to an offline dataset collected by a behavior policy  $\pi_{\beta}$ . Using this supplementary dataset, the expert matching objective of DemoDICE is instantiated over ValueDICE:
$$
- D _ {K L} \left(d ^ {\pi} (s, a) \| d ^ {\pi_ {E}} (s, a)\right) - \alpha D _ {K L} \left(d ^ {\pi} (s, a) \| d ^ {\pi_ {\beta}} (s, a)\right),
$$
where  $\alpha$  is a positive weight for the constraint.
The above optimization objective can be transformed into three tractable components: 1) a reward function  $r(s, a)$  derived by pre-training a binary discriminator  $D: S \times \mathcal{A} \rightarrow [0,1]$ :
$$
r (s, a) = - \log (\frac {1}{D ^ {*} (s , a)} - 1),
$$
$$
D ^ {*} (s, a) = \underset {D} {\arg \max } = \mathbb {E} _ {d ^ {\pi_ {E}}} [ \log D (s, a) ] + \mathbb {E} _ {d ^ {\pi_ {\beta}}} [ \log (1 - D (s, a)) ],
$$
2) a value function optimization objective:
$$
\mathcal {J} (v) = - (1 - \gamma) \mathbb {E} _ {s \sim \rho_ {0}} [ v (s) ] - (1 + \alpha) \log \mathbb {E} _ {(s, a) \sim d ^ {\pi_ {\beta}}} \left[ \exp (\frac {r (s , a) + \mathbb {E} _ {s ^ {\prime} \sim P (s , a)} (v (s ^ {\prime})) - v (s)}{1 + \alpha}) \right],
$$
and 3) a policy optimization step:
$$
\mathcal {J} (\pi) = \mathbb {E} _ {(s, a) \sim d ^ {\pi_ {\beta}}} \left[ v ^ {*} (s, a) \log \pi (a | s) \right],
$$
$$
v ^ {*} (s, a) = \underset {v} {\arg \max } \mathcal {J} (v).
$$
We report the offline results using the official Tensorflow implementation<sup>13</sup>.
$^{12}$ https://github.com/google-research/google-research/tree/master/value_dice
$^{13}$ https://github.com/KAIST-AILab/imitation-dice
State Matching Offline Distribution Correction Estimation (SMODICE). SMODICE [56] proposes to solve offline IL tasks in LfO and cross-domain settings and it optimizes the following state occupancy objective:
$$
- D _ {K L} \left(d ^ {\pi} (s) \| d ^ {\pi_ {E}} (s)\right).
$$
To incorporate the offline dataset, SMODICE derives an f-divergence regularized state-occupancy objective:
$$
\mathbb {E} _ {s \sim d ^ {\pi} (s)} \left[ \log \left(\frac {d ^ {\pi_ {\beta}} (s)}{d ^ {\pi_ {E}} (s)}\right) \right] + - D _ {f} \left(d ^ {\pi} (s, a) \| d ^ {\pi_ {\beta}} (s, a)\right).
$$
Intuitively, the first term can be interpreted as matching the offline states towards the expert states, while the second regularization term constrains the policy close to the offline distribution of state-action occupancy. Similarly, we can divide the objective into three steps: 1) deriving a state-based reward by learning a state-based discriminator:
$$
\begin{array}{l} r (s, a) = - \log (\frac {1}{D ^ {*} (s)} - 1), \\ D ^ {*} (s, a) = \underset {D} {\arg \max } = \mathbb {E} _ {d ^ {\pi_ {E}}} [ \log D (s) ] + \mathbb {E} _ {d ^ {\pi_ {\beta}}} [ \log (1 - D (s)) ], \\ \end{array}
$$
2) learning a value function using the learned reward:
$$
\mathcal {J} (v) = - (1 - \gamma) \mathbb {E} _ {s \sim \rho_ {0}} [ v (s) ] - \log \mathbb {E} _ {(s, a) \sim d ^ {\pi_ {\beta}}} \left[ f _ {*} (r (s, a) + \mathbb {E} _ {s ^ {\prime} \sim P (s, a)} (v (s ^ {\prime})) - v (s)) \right],
$$
and 3) training the policy via weighted regression:
$$
\begin{array}{l} \mathcal {J} (\pi) = \mathbb {E} _ {(s, a) \sim d ^ {\pi_ {\beta}}} \left[ f _ {*} ^ {\prime} (r (s, a) + \mathbb {E} _ {s ^ {\prime} \sim P (s, a)} \left(v ^ {*} \left(s ^ {\prime}\right)\right) - v ^ {*} (s)\right) \log \pi (a | s) \bigg ], \\ v^{*}(s,a) = \operatorname *{arg  max}_{v}\mathcal{J}(v), \\ \end{array}
$$
where  $f_{*}$  is the Fenchel conjugate of f-divergence (please refer to Ma et al. [56] for more details).
We conduct experiments using the official Pytorch implementation  $^{14}$ , where the f-divergence used is  $\mathcal{X}^2$ -divergence. On the LfD tasks, we change the input of the discriminator to state-action pairs.
$^{14}$ https://github.com/JasonMa2016/SMODICE