SlowGuess/NeurIPS_2023_OCR huggingface dataset:
where  $\lambda = (\alpha -\beta) / \alpha$  and  $\mathbb{E}_{\boldsymbol{s},\boldsymbol{a}\sim \mathcal{D}}[\log p(\boldsymbol {a}|\boldsymbol {s})]$  is constant. The regression objective bounds the IRL objective for a worst-case on-policy state distribution  $\mu_q(\pmb {s})$ , which motivates its scaling through  $\beta$ . If  $\mathcal{D}$  has sufficient coverage, no RL finetuning or KL regularization is required, so  $\beta = \alpha$  and  $\lambda = 0$ . If  $\mathcal{D}$  does not have sufficient coverage, then let  $\beta < \alpha$  so  $\lambda >0$  to regularize the BC fit and finetune the policy with RL accordingly with additional soft policy iteration steps.
Task	Random return	Expert return
HalfCheetah-v2	-282	8770
Hopper-v2	18	2798
Walker2d-v2	1.6	4118
Ant-v2	-59	5637
Humanoid-v2	123	9115
door-expert-v0	-56	2882
hammer-expert-v0	-274	12794
Table 2: Expert and random policy returns used to normalize the performance for Gym and Adroit tasks. Taken from Orsini et al. [11].
Proof. We begin with the KL-regularized game-theoretic IRL objective (c.f. Equation 1 [17]),
$$
\mathcal {J} (r, \pi , \beta) = \mathbb {E} _ {\boldsymbol {s}, \boldsymbol {a} \sim \mathcal {D}} \left[ r (\boldsymbol {s}, \boldsymbol {a}) \right] - \mathbb {E} _ {\boldsymbol {s}, \boldsymbol {a} \sim \rho_ {\pi}} \left[ r (\boldsymbol {s}, \boldsymbol {a}) - \beta \log \frac {\pi (\boldsymbol {a} \mid \boldsymbol {s})}{p (\boldsymbol {a} \mid \boldsymbol {s})} \right].
$$
Substituting the coherent reward from Theorem 1, the objective becomes
$$
\mathcal {J} (\pmb {\theta}, \pi , \beta) = \mathbb {E} _ {\pmb {s}, \pmb {a} \sim \mathcal {D}} \left[ \alpha \log \frac {q _ {\pmb {\theta}} (\pmb {a} \mid \pmb {s})}{p (\pmb {a} \mid \pmb {s})} \right] - \mathbb {E} _ {\pmb {s}, \pmb {a} \sim \rho_ {\pi}} \left[ \alpha \log \frac {q _ {\pmb {\theta}} (\pmb {a} \mid \pmb {s})}{p (\pmb {a} \mid \pmb {s})} - \beta \log \frac {\pi (\pmb {a} \mid \pmb {s})}{p (\pmb {a} \mid \pmb {s})} \right].
$$
Due to the realizability assumption (Assumption 1), the optimal policy  $\pi_{\theta_*}$  can be obtained through maximum likelihood estimation  $\pmb{\theta}_{*} = \arg \max_{\pmb{\theta}}\mathbb{E}_{\pmb{s},\pmb{a}\sim \mathcal{D}}[\log q_{\pmb{\theta}}(\pmb {s},\pmb {a})]$ . With this assumption, the left-hand term performs maximum likelihood estimation of the optimal policy. Therefore, combined with the coherent reward which is derived from policy inversion, we replace the inner policy optimization of  $\pi$  with the maximum likelihood solution of  $q_{\theta}$ , resulting in a simpler objective  $\tilde{\mathcal{J}}$  in terms of only  $\pmb{\theta}$ , which matches the objective  $\mathcal{J}$  exactly when  $\beta = \alpha$  and  $\pi = q_{\theta}$  (Lemma 2). We consider  $\tilde{\mathcal{J}}$  with a small negative perturbation  $\beta = \alpha -\delta \alpha ,\delta \alpha \geq 0$  to examine the effect of the right-term term,
$$
\tilde {\mathcal {J}} (\pmb {\theta}) = \mathcal {J} (\pmb {\theta}, q _ {\pmb {\theta}}, \alpha - \delta \alpha) = \mathbb {E} _ {\pmb {s}, \pmb {a} \sim \mathcal {D}} \left[ \alpha \log \frac {q _ {\pmb {\theta}} (\pmb {a} \mid \pmb {s})}{p (\pmb {a} \mid \pmb {s})} \right] - \mathbb {E} _ {\pmb {a} \sim q _ {\pmb {\theta}} (\cdot | \pmb {s}), \pmb {s} \sim \mu_ {q _ {\pmb {\theta}}}} \left[ \delta \alpha \log \frac {q _ {\pmb {\theta}} (\pmb {a} \mid \pmb {s})}{p (\pmb {a} \mid \pmb {s})} \right],
$$
The right-hand term contains a conditional KL divergence, which we can replace with an upper-bound using the data processing inequality from Lemma 3 and the weight-space KL. This weight-space KL is constant w.r.t. the state distribution expectation and we introduce the scale factor  $\lambda = \delta \alpha / \alpha$ . The resulting lower-bound objective is proportional to the regularized behavioral cloning objective proposed in Equation 8,
$$
\tilde {\mathcal {J}} (\pmb {\theta}) \geq \alpha (\mathbb {E} _ {\pmb {s}, \pmb {a} \sim \mathcal {D}} [ \log q _ {\pmb {\theta}} (\pmb {a} \mid \pmb {s}) ] - \lambda \mathbb {D} _ {\mathrm {K L}} [ q _ {\pmb {\theta}} (\pmb {w}) | | p (\pmb {w}) ] + C), \quad C = \mathbb {E} _ {\pmb {s}, \pmb {a} \sim \mathcal {D}} [ - \log p (\pmb {a} \mid \pmb {s}) ].
$$
The intuition from Theorem 2 is that the KL regularization during BC achieves reward shaping similar to the IRL saddle-point solution when using the coherent reward. The lower bound's 'tightness' depends on the data processing inequality and is difficult to reason about. Therefore, this result demonstrates that KL-regularized BC provides useful reward shaping from the game-theoretic perspective, rather than provide a tight lower bound. To reinforce this result, we visualize the effect in Appendix G.
L Implementation and experimental details
Our codebase for this work is open-sourced and is available at github.com/google-deepmind/csil.
The tabular experiments are implemented in jax automatic differentiation and linear algebra library [102]. The continuous state-action experiments are implemented in acne, again using jax and its ecosystem [103]. One feature of acne is distributed training, where we use four concurrent actors to speed up the wall clock learning time. IQLearn and PPIL were re-implemented based on the open-source implementations as reference [100, 104]. CSIL, IQLearn, and PPIL shared a combined
implementation as 'soft' imitation learning algorithms due to their common structure and to facilitate ablation studies. Baselines DAC, SQIL, PWIL, SACfd, and CQL are implemented in acne. We used a stochastic actor for learning and a deterministic actor (i.e., evaluating the mean action) for evaluation.
The DemoDICE and SMODICE offline imitation learning baselines were implemented using the original implementations [105, 106]. This meant that the expert demonstrations were obtained from d4r1 rather than Orsini et al. [11], but the returns were normalized to the d4r1 returns for consistency. We note that there is a non-negligible discrepancy between the d4r1 and Orsini et al. [11] return values for the MuJoCo tasks, but in the imitation setting, comparing normalized values should be a meaningful metric.
Due to the common use of SAC as the base RL algorithm, most implementations shared common standard hyperparameters. The policy and critic networks were comprised of two layers with 256 units and ELU activations. Learning rates were 3e-4, the batch size was 256, and the target network smoothing coefficient was 0.005. One difference to standard RL implementations is the use of layernorm in the critic, which is adopted in some algorithm implementations in acne and was found to improve performance. For the inverse temperature  $\alpha$ , for most methods, we used the adaptive strategy based on a lower bound on the policies entropy. IQLearn and PPIL both have a fixed temperature strategy that we tuned in the range [1.0, 0.3, 0.1, 0.03, 0.01]. CSIL also has a fixed temperature strategy but is applied to the KL penalty. In practice, they are very similar, so we swept the same range. We also investigated a hard KL bound and adaptive temperature strategy, but there was no performance improvement to warrant the extra complexity. CSIL's HETSTAT policies had an additional bottleneck layer of 12 units and a final layer of 256 units with stationary activations. For the Gym tasks, triangular activations were used. For robomimic tasks, periodic ReLU activations were used. IQLearn, PPIL, and CSIL all have single critics rather than dual critics, while the other baselines had dual critics to match the standard SAC implementation. We tried dual critics to stabilize IQLearn and PPIL, but this did not help. CSIL's reward finetuning learning rate was 1e-3. As the reward model was only being finetuned with non-expert samples, we did not find this hyperparameter to be very sensitive. For robomimic, larger models with two layers of 1024 units were required. The HETSTAT policies increased accordingly with a 1024 unit final layer and 48 units for the bottleneck layer.
For DAC, we used the hyperparameters and regularizers recommended by Orsini et al. [11], including spectral normalization for the single hidden layer discriminator with 64 units. IQLearn has a lower actor learning rate of 3e-5, and the open-source implementation has the option of different value function regularizers per task [100]. We used 'v0' regularization for all environments, which minimizes the averaged value function at the expert samples. In the original IQLearn experiments, this regularizer is tuned per environment [100]. PPIL shares the same smaller actor learning rate of 3e-5 and the open-source implementation also uses 20 critic updates for every actor update, which we maintained. The implementation of PPIL borrows the same value function regularization as IQLearn, so we used 'v0' regularization for all environments [104]. For SACfD, we swept the demo-to-replay ratio across [0.1, 0.01]. For Adroit, the ratio was 0.01, and for robomimic it was 0.1.
The major implementation details of CSIL are discussed in Sections 4, C, D and E. The pre-training hyperparameters are listed in Table 4. While the critic pre-training is not too important for performance (Figure 44), appropriate policy pre-training is crucial. Instead of training a regularization
Experiment	Algorithm
DAC	PWIL	SQIL	IQLearn	PPIL	CSIL
Online Gym	H(π)	H(π)	H(π)	0.01	0.01	0.01
Offline Gym	H(π)	H(π)	H(π)	0.03	0.03	0.1
Online Humanoid-v2	H(π)	H(π)	H(π)	0.01	0.01	0.01
Online Adroit	H(π)	H(π)	H(π)	0.03	0.03	0.1
Online robomimic	H(π)	-	-	0.1	-	0.1
Offline robomimic	-	-	-	-	-	1.0
Image-based robomimic	-	-	-	-	-	0.3
Table 3: The entropy regularization hyperparameter value used across methods and environments. Values refer to a fixed temperature while  $\mathbb{H}(\pi)$  refers to the adaptive temperature scheme, which ensures the policy entropy is at least equal to  $-d_{a}$ , where  $d_{a}$  is the dimension of the action space.
Experiment	Pre-training hyperparameters
n(π)	λ(π)	n(Q)	λ(Q)
Online and offline Gym	25000	1e-3	5000	1e-3
Online Humanoid-v2	500	1e-3	5000	1e-3
Online Adroit	25000	1e-4	5000	1e-3
Online and offline robomimic	50000	1e-4	2000	1e-3
Image-based robomimic	25000	1e-4	2000	1e-3
Table 4: CSIL's pre-training hyperparameters for the policy  $(\pi)$  and critic  $(Q)$ , listing the learning rate  $(\lambda)$  and number of iterations  $(n)$ .
hyperparameter, we opted for early stopping. As a result, there is a trade-off between BC accuracy and policy entropy. In the online setting, if the policy is trained too long, the entropy is reduced too far, and there is not enough exploration for RL. As a result, the policy pre-training hyperparameters need tuning when action dimension and dataset size change significantly to ensure the BC fit is sufficient for downstream performance. For robomimic, we used negative CSIL rewards scaled by 0.5 to ensure they lie around  $[-1,0]$ . For vision-based robomimic, the CNN torso was an impala-style [107] Acme ResNetTorso down to a LayerNorm MLP with 48 output units, ELU activations, orthogonal weight initialization, and final tanh activation. Mandlekar et al. [84] used a larger ResNet18. As in Mandlekar et al. [84], image augmentation was achieved with random crops from 84 down to 76. Pixel values were normalized to  $[-0.5,0.5]$ . For memory reasons, the batch size was reduced to 128. The implementations of IQLearn, PPIL and CSIL combined equally-sized mini-batches from the expert and non-expert replay buffers for learning the reward, critic and policy. We did not investigate tuning this ratio.
Regarding computational resources, the main experiments were run using acme, which implements distributed learning. As a result, our 'learner' (policy evaluation and improvement) runs on a single TPU v2. We ran four actors to interact with the environment. Depending on the algorithm, there were also one or more evaluators. For vision-based tasks, we used A100 GPUs for the vision-based policies.
M Experimental results
This section includes additional experimental results from Section 5. Videos of the simulation studies are available at joemwatson.github.io/csil.
M.1 Inverse optimal control
This section shows the converged value function, stationary distribution, and 'windy' stationary distribution for the tabular setting results from Table 1. 'Oracle classifier' constructs an offline classifier-based reward separating the states and actions with- and without demonstration data, in the spirit of methods such as SQIL and ORIL [85]. No optimization was conducted on the windy environment, just evaluation to assess the policy robustness to disturbance. We use the Viridis colour map, so yellow regions have a larger value than blue regions.
Figure 12: Expert on the dense MDP.
Figure 14: Oracle classifier on the dense MDP.
Figure 16: IQLearn on the dense MDP.
Figure 18: ME-IRL on the dense MDP.
Figure 13: BC on the dense MDP.
Figure 15: GAIL on the dense MDP.
Figure 17: PPIL on the dense MDP.
Figure 19: CSIL on the dense MDP.
Value function
Figure 20: Expert on the sparse MDP.
Nominal stationary distribution
Windy stationary distribution
Value function
Figure 22: Oracle classifier on the sparse MDP.
Value function
Figure 24: IQLearn on the sparse MDP.
Value function
Figure 26: ME-IRL on the sparse MDP.
Nominal stationary distribution
Figure 21: BC on the sparse MDP.
Value function
Figure 23: GAIL on the sparse MDP.
Value function
Figure 25: PPIL on the sparse MDP.
Value function
Figure 27: CSIL on the sparse MDP.
Figure 28: Normalized performance of CSIL again baselines for online imitation learning on MuJoCo Gym tasks. Uncertainty intervals depict quartiles over ten seeds. CSIL exhibits stable convergence with both sample and demonstration efficiency. BC here is applied to the demonstration dataset.  $n$  refers to demonstration trajectories.
Figure 29: Normalized performance of CSIL against baselines for online imitation learning for Adroit tasks. Uncertainty intervals depict quartiles over ten seeds. CSIL exhibits stable convergence with both sample and demonstration efficiency. Many baselines cannot achieve stable convergence due to the high-dimensional action space. SACfd is an oracle baseline that combines the demonstrations with the true (shaped) reward.  $n$  refers to demonstration trajectories.
Figure 30: Normalized performance of CSIL against baselines for offline imitation learning for Gym tasks. Uncertainty intervals depict quartiles over ten seeds. While all methods struggle in this setting compared to online learning, CSIL manages convergence with enough demonstrations. CQL is an oracle baseline using the true rewards. The BC baseline is trained on the whole offline dataset, not just demonstrations.  $n$  refers to demonstration trajectories.
Figure 31: Normalized performance of CSIL against baselines for online imitation learning for Humanoid-v2. Uncertainty intervals depict quartiles over ten seeds.  $n$  refers to demonstration trajectories.
M.2 Continuous control from agent demonstrations
Step-wise results. Figure 28, 29 and 30 show step-wise performance curves, as opposed to the demonstration-wise results in the main paper. For online experiments, steps refers to the actor and learner, while for offline experiments, steps correspond to the learner.
Humanoid-v2 results. We evaluated CSIL and baselines on the high-dimensional Humanoid-v2 locomotion task in Figure 31. We found that CSIL struggled to solve this task due to the difficulty for BC to perform well at this task, even when using all 200 of the available demonstrations in the dataset.
One aspect of Humanoid-v2 to note is that a return of around 5000 (around 0.5 normalized) corresponds to the survival / standing bonus, rather than actual locomotion (around 9000). While the positive nature of the coherent reward enabled learning the stabilize, this was not consistent across seeds. We also found it difficult to reproduce baseline results due to IQLearn's instability, which we also reproduced on the author's implementation [100]. In theory, BC should be able to approximate the original optimal agent's policy with sufficient demonstrations. We believe the solution may be additional BC regularization, as the environment has a 376-dimensional state space.
Performance improvement of coherent imitation learning A desired quality of CSIL is its ability to match or exceed its initial BC policy. Figures 32, 33 and 34 show the performance of CSIL relative to its initial BC policy. Empirical, CSIL matches or exceeds the initial BC policy across environments.
Figure 32: Normalized performance of CSIL for online Gym tasks against its BC initialization. Uncertainty intervals depict quartiles over ten seeds.  $n$  refers to demonstration trajectories.
Figure 33: Normalized performance of CSIL for online Adroit tasks against its BC initialization. Uncertainty intervals depict quartiles over ten seeds.  $n$  refers to demonstration trajectories.
Figure 34: Normalized performance of CSIL for offline Gym tasks against its BC initialization. Uncertainty intervals depict quartiles over ten seeds.  $n$  refers to demonstration trajectories.
Figure 35: We investigate the weakness of IQLearn and PPIL for offline learning (Ant-v2) and high-dimensional action spaces (door-v0). Compared to CSIL, IQLearn and PPIL have less stable and effective  $Q$  values, in particular for the expert demonstrations, which results in the policy not inaccurately predicting the demonstration actions (expert_mse).
Divergence in minimax soft imitation learning To investigate the performance issues in IQLearn and PPIL relative to CSIL, we monitor a few key diagnostics during training for online door-v0 and offline Ant-v2 in Figure 35. One key observation is the drift in the BC accuracy. We attribute this to the instability in the reward or critic due to the saddle-point optimization.
M.3 Continuous control from human demonstrations
Figure 36 shows performance on the online setting w.r.t. learning steps, and Figure 37 shows the offline performance w.r.t. learning steps. For online learning, CSIL is consistently more sample efficient than the oracle baseline SACfD, presumably due to BC pre-training and CSIL's shaped reward w.r.t. robomimic's sparse reward on success. CSIL was also able to surpass Mandlekar et al. [84]'s success rate on NutAssemblySquare. For offline learning from suboptimal human demonstrations, it appears to be a much harder setting for CSIL to improve on the initial BC policy, especially for harder tasks.
Figure 36: Average success rate over 50 evaluations for online imitation learning for robomimic tasks. Uncertainty intervals depict quartiles over ten seeds. The BC policy is trained on the demonstration data.  $n$  refers to demonstration trajectories.
Figure 37: Average success rate over 50 evaluations for offline imitation learning for robomimic tasks. Uncertainty intervals depict quartiles over ten seeds. This experiment shows that learning from suboptimal offline demonstrations is much harder than on-policy data. It is also harder to use the KL regularization to ensure improvement over the initial policy.  $n$  refers to demonstration trajectories.
N Ablation studies
This section includes ablation studies for the CSIL algorithm and baselines.
N.1 Baselines with behavioral cloning pre-training
To assess the importance of BC pre-training, we evaluate the performance of IQLearn and PPIL with BC pre-training in Figure 38. While difficult to see across all experiments, there is a small drop in performance at the beginning of training as the policy 'unlearns' as the critic and reward are randomly initialized since there is no notion of coherency.
N.2 Baselines with behavioral cloning pre-training and KL regularization
To bring IQLearn and PPIL closer to CSIL, we implement the baselines with BC pre-training and KL regularization, but with SAC-style policies rather than HETSTAT. We evaluate on the harder tasks, online Adroit, and offline Gym, to see if these features contribute sufficient stability to solve the task. Both Figure 39 and Figure 40 show an increase in initial performance, it is not enough to completely stabilize the algorithms.
Figure 38: Normalized performance of IQLearn and PPIL for online Gym tasks with BC pre-training. Uncertainty intervals depict quartiles over ten seeds.  $n$  refers to demonstration trajectories.
Figure 39: Normalized performance of IQLearn and PPIL for online Adroit tasks with BC pretraining. Uncertainty intervals depict quartiles over ten seeds.  $n$  refers to demonstration trajectories.
Figure 40: Normalized performance of IQLearn and PPIL for offline Gym tasks with BC pre-training. Uncertainty intervals depict quartiles over ten seeds.  $n$  refers to demonstration trajectories.
Figure 41: Normalized performance of IQLearn and PPI for online Gym tasks with HETSTAT policies. Uncertainty intervals depict quartiles over ten seeds.  $n$  refers to demonstration trajectories.
N.3 Baselines with heteroscedastic stationary policies
To investigate the impact of HETSTAT policies alone, we evaluated PPIL and IQLearn with HETSTAT policies on the online Gym tasks in Figure 41. There is not a significant discrepancy between MLP and HETSTAT policies.
N.4 Baselines with pre-trained heteroscedastic stationary policies as priors
To bring IQLearn and PPIL closer to CSIL, we implement the baselines with BC pre-training, HETSTAT policies and KL regularization with stationary policies. We evaluate on the harder tasks, online Adroit and offline Gym, to see if these features contribute sufficient stability to solve the task. Both Figure 42 and Figure 43 show an increase in performance, especially with increasing number of demonstrations, but it is ultimately not enough to completely stabilize the algorithms to the degree seen in CSIL.
Figure 42: Normalized performance of IQLearn and PPIL for online Adroit tasks with pre-trained HETSTAT policies and KL-regularization. Uncertainty intervals depict quartiles over ten seeds.  $n$  refers to demonstration trajectories.
Steps  $(1\times 10^{6})$
Figure 43: Normalized performance of IQLearn and PPIL for offline Gym tasks with pre-trained HETSTAT policies and KL-regularization. Uncertainty intervals depict quartiles over ten seeds.  $n$  refers to demonstration trajectories.
Figure 44: Normalized performance of CSIL for online Gym tasks a critic pre-training ablation. Uncertainty intervals depict quartiles over ten seeds.  $n$  refers to demonstration trajectories.
Figure 45: Normalized performance of CSIL for online Gym tasks with a reward finetuning ablation. Uncertainty intervals depict quartiles over ten seeds. Refinement is most important in the low demonstration setting where the initial BC policy and coherent reward is less defined, and therefore the initial policy is poor and the coherent reward is more susceptible to exploitation of the stationary approximation errors.  $n$  refers to demonstration trajectories.
N.5 Importance of critic pre-training
To ablate the importance of critic pre-training, Figure 44 shows the performance difference on online Gym tasks. Critic pre-training provides only a minor performance improvement.
N.6 Importance of reward finetuning
To ablate the importance of reward finetuning, Figure 45. shows the performance difference on online Gym tasks. Reward finetuning is crucial in the setting of few demonstration trajectories, as the imitation reward is sparse and therefore csil is more susceptible to approximation errors in stationarity. It is also important for reducing performance variance and ensuring stable convergence.
N.7 Importance of critic regularization
To ablate the importance of critic Jacobian regularization, Figure 46. shows the performance difference on online Gym tasks. The regularization has several benefits. One is faster convergence, as it encodes the coherency inductive bias into the critic. A second is stability, as it encodes first-order optimality w.r.t. the expert demonstrations, mitigating errors in policy evaluation. The inductive bias also appears to reduce performance variance. One observation is that the performance of Ant-v2 improves without the regularization. This difference could be because the regularization encourages imitation rather than apprenticeship-style learning where the agent improves on the expert.