SlowGuess/NeurIPS_2023_OCR huggingface dataset:
These desiderata are straightforward to achieve for maximum entropy priors in the tabular setting, where count-based conditional density estimation is combined with the prior, e.g.,  $p(\boldsymbol{a} \mid \boldsymbol{s}) = \mathcal{U}_{\mathcal{A}}(\boldsymbol{a})$ . Unfortunately, point 3 is challenging in the continuous setting when using function approximation, as shown in Figure 3. However, we can adopt ideas from Gaussian process theory to approximate such policies, which is discussed further in Section 4. The learning objective combines maximizing the likelihood of the demonstrations w.r.t. the predictive distribution, as well as KL regularization against the prior weight distribution, performing regularized heteroscedastic regression [63],
$$
\max  _ {\boldsymbol {\theta}} \mathbb {E} _ {\boldsymbol {a}, \boldsymbol {s} \sim \mathcal {D}} [ \log q _ {\boldsymbol {\theta}} (\boldsymbol {a} \mid \boldsymbol {s}) ] - \lambda \mathbb {D} _ {\mathrm {K L}} [ q _ {\boldsymbol {\theta}} (\boldsymbol {w}) | | p (\boldsymbol {w}) ], q _ {\boldsymbol {\theta}} (\boldsymbol {a} \mid \boldsymbol {s}) = \int p (\boldsymbol {a} \mid \boldsymbol {s}, \boldsymbol {w}) q _ {\boldsymbol {\theta}} (\boldsymbol {w}) d \boldsymbol {w}. \tag {8}
$$
This objective has been used to fit parametric Gaussian processes with a similar motivation [64].
We now describe how this BC approach and the coherent reward are used for imitation learning (Algorithm 1) when combined with soft policy iteration algorithms, such as SAC.
On a high level, Algorithm 1 can be summarized by the following steps,
1. Perform regularized BC on the demonstrations with a parametric stochastic policy  $q_{\theta}(a \mid s)$ .
2. Define the coherent reward,  $r_{\theta}(s,a) = \alpha (\log q_{\theta}(a\mid s) - \log p(a\mid s))$  , with temperature  $\alpha$
3. With temperature  $\beta < \alpha$ , perform SPI to improve on the BC policy with the coherent reward.
The reduced temperature  $\beta$  is required to improve the policy after inverting it with temperature  $\alpha$ .
This coherent approach contrasts combining BC and prior actor-critic IRL methods, where learning the reward and critic from scratch can lead to 'unlearning' the initial BC policy [11]. Moreover, while our method does not seek to learn the true underlying reward, but instead opts for one shaped by the environment used to generate demonstrations, we believe this is a reasonable compromise in practice. Firstly, CSIL uses IRL to tackle the covariate shift problem in BC, as the reward shaping (Definition 1) encourages returning to the demonstration distribution when outside it, which is also the goal of prior imitation methods, e.g., [23, 24]. Secondly, in the entropy-regularized setting, an IRL method has the drawback of requiring demonstrations generated from two different environments or discount factors to accurately infer the true reward [54], which is often not readily available in practice, e.g., when teaching a single robot a task. Finally, additional MLPs could be used to estimate the true (unshaped) reward from data [65, 53] if desired. We also provide a divergence minimization perspective of CSIL in Section J of the Appendix. Another quality of CSIL is that it is conceptually simpler than alternative IRL approaches such as AIRL. As we simply use a shaped estimate of the true reward with entropy-regularized RL, CSIL inherits the theoretical properties of these regularized algorithms [17, 66] regarding performance. As a consequence, this means that analysis of the imitation quality requires analyzing the initial behavioral cloning procedure.
Figure 3: For the coherent reward to be effective, the stochastic policy should return to the prior distribution outside of the data distribution (left). The typical heteroscedastic parametric policies have undefined out-of-distribution behavior and typically collapse to the action limits due to the network extrapolation and tanh transformation (middle). By approximating stationary Gaussian processes, we can design policies that exhibit the desired behavior with minimal network modifications (right).
Partial expert coverage, minimax optimization, and regularized regression. The reward shaping theory from Lemma 1 shows that an advantage function can be used as a shaped reward function when the state potential is a value function. By inverting the soft policy update, the log policy ratio acts as a form of its advantage function. However, in practice, we do not have access to the complete advantage function but rather an estimate due to finite samples and partial state coverage. Suppose the expert demonstrations provide only partial coverage of the state-action space. In this case, the role of an inverse reinforcement learning algorithm is to use additional knowledge of the MDP, e.g., online samples, a dynamics model, or an offline dataset, to improve the reward estimate. As CSIL uses the log policy ratio learned only from demonstration data, how can it be a useful shaped reward estimate? We show in Theorem 2 that using entropy regularization in the initial behavioral cloning step plays a similar role to the saddle-point optimization in game-theoretic IRL.
Theorem 2. (Coherent inverse reinforcement learning as KL-regularized behavioral cloning). A KL-regularized game-theoretic IRL objective, with policy  $q_{\theta}(\boldsymbol{a} \mid \boldsymbol{s})$  and coherent reward parameterization  $r_{\theta}(\boldsymbol{s}, \boldsymbol{a}) = \alpha (\log q_{\theta}(\boldsymbol{a} \mid \boldsymbol{s}) - \log p(\boldsymbol{a} \mid \boldsymbol{s}))$  where  $\alpha \geq 0$ , is lower bounded by a scaled KL-regularized behavioral cloning objective and a constant term when the optimal policy, posterior, and prior share a hypothesis space  $p(\boldsymbol{a} \mid \boldsymbol{s}, \boldsymbol{w})$  and finite parameters  $\boldsymbol{w}$ ,
$$
\begin{array}{l} \mathbb {E} _ {\boldsymbol {s}, \boldsymbol {a} \sim \mathcal {D}} \left[ r _ {\boldsymbol {\theta}} (\boldsymbol {s}, \boldsymbol {a}) \right] - \mathbb {E} _ {\boldsymbol {a} \sim q _ {\boldsymbol {\theta}} (\cdot | \boldsymbol {s}), \boldsymbol {s} \sim \mu_ {q _ {\boldsymbol {\theta}}} (\cdot)} \left[ r _ {\boldsymbol {\theta}} (\boldsymbol {s}, \boldsymbol {a}) - \beta \left(\log q _ {\boldsymbol {\theta}} (\boldsymbol {a} \mid \boldsymbol {s}) - \log p (\boldsymbol {a} \mid \boldsymbol {s})\right) \right] \geq \\ \alpha \left(\mathbb {E} _ {\boldsymbol {s}, \boldsymbol {a} \sim \mathcal {D}} [ \log q _ {\boldsymbol {\theta}} (\boldsymbol {a} \mid \boldsymbol {s}) ] - \lambda \mathbb {D} _ {\mathrm {K L}} [ q _ {\boldsymbol {\theta}} (\boldsymbol {w}) | | p (\boldsymbol {w}) ] + \mathbb {E} _ {\boldsymbol {s}, \boldsymbol {a} \sim \mathcal {D}} [ \log p (\boldsymbol {a} \mid \boldsymbol {s}) ]\right). \\ \end{array}
$$
where  $\lambda = (\alpha -\beta) / \alpha$  and  $\mathbb{E}_{s,a\sim \mathcal{D}}[\log p(\boldsymbol {a}\mid 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.
The proof is provided in Appendix K.3. In the tabular setting, it is straightforward for the cloned policy to reflect the prior distribution outside of the expert's data distribution. However, this behavior is harder to capture in the continuous setting where function approximation is typically adopted.
4 Stationary processes for continuous control policies
Maximum entropy methods used in reinforcement learning can be recast as a minimum relative entropy problem against a regularizing prior policy with a uniform action distribution, i.e., a prior  $p(\boldsymbol{a} \mid \boldsymbol{s}) = \mathcal{U}_{\mathcal{A}}(\boldsymbol{a}) \forall \boldsymbol{s} \in \mathcal{S}$ , where  $\mathcal{U}_{\mathcal{A}}$  is the uniform distribution over  $\mathcal{A}$ . Achieving such a policy in the tabular setting is straightforward, as the policy can be updated independently for each state. However, such a policy construction is far more difficult for continuous states due to function approximation capturing correlations between states. To achieve the desired behavior, we can use stationary process theory (Definition 3) to construct an appropriate function space (Figure 3).
Definition 3. (Stationary process, Cox and Miller [67]). A process  $f: \mathcal{X} \to \mathcal{Y}$  is stationary if its joint distribution in  $\pmb{y} \in \mathcal{Y}$ ,  $p(\pmb{y}_1, \dots, \pmb{y}_n)$ , is constant w.r.t.  $\pmb{x}_1, \dots, \pmb{x}_n \in \mathcal{X}$  and all  $n \in \mathbb{N}_{>0}$ .
To approximate a stationary policy using function approximation, Gaussian process (GP) theory provides a means using features in a Gaussian linear model that defines a stationary process [62]. To approximate such feature spaces using neural networks, this can be achieved through a relatively wide final layer with a periodic activation function  $(f_{\mathrm{per}})$ , which can be shown to satisfy the stationarity property [68]. Refer to Appendix C for technical details. To reconcile this approach with prior work such as SAC, we use the predictive distribution of our Gaussian process in lieu of a network directly
predicting Gaussian moments. The policy is defined as  $\pmb{a} = \tanh(z(\pmb{s}))$ , where  $\pmb{z}(\pmb{s}) = \pmb{W}\phi(\pmb{s})$ . The weights are factorized row-wise  $\pmb{W} = [\pmb{w}_1, \dots, \pmb{w}_{d_a}]^\top$ ,  $\pmb{w}_i = \mathcal{N}(\pmb{\mu}_i, \pmb{\Sigma}_i)$  to define a GP with independent actions. Using change-of-variables like SAC [17], the policy is expressed per-action as
$$
q (a _ {i} \mid s) = \mathcal {N} \left(z _ {i}; \boldsymbol {\mu} _ {i} ^ {\top} \phi (s), \phi (s) ^ {\top} \boldsymbol {\Sigma} _ {i} \phi (s)\right) \cdot \left| \det  \left(\frac {\mathrm {d} a _ {i}}{\mathrm {d} z _ {i}}\right) \right| ^ {- 1}, \quad \phi (s) = f _ {\text {p e r}} (\tilde {\boldsymbol {W}} \phi_ {\text {m l p}} (s)).
$$
$\tilde{W}$  are weights drawn from a distribution, e.g., a Gaussian, that also characterizes the stochastic process and  $\phi_{\mathrm{mlp}}$  is an arbitrary MLP. We refer to this heteroscedastic stationary model as HETSTAT.
Function approximation necessitates several additional practical implementation details of CSIL.
Approximating and regularizing the critic. Theorem 1 and Lemma 1 show that the log policy ratio is also a shaped critic. However, we found this model was not expressive enough for further policy evaluation. Instead, the critic is approximated using as a feedforward network and pre-trained after the policy using SARSA [8] and the squared Bellman error. For coherency, a useful inductive bias is to minimize the critic Jacobian w.r.t. the expert actions to approximate first-order optimality, i.e.,  $\min_{\phi} \mathbb{E}_{s,a \sim \mathcal{D}}[\nabla_a Q_\phi(s,a)]$ , as an auxiliary critic loss. For further discussion, see Section H in the Appendix. Pre-training and regularization are ablated in Figures 44 and 46 in the Appendix.
Using the cloned policy as prior. While a uniform prior is used in the CSIL reward throughout, in practice, we found that finetuning the cloned policy with this prior in the soft Bellman equation (Equation 3) leads to divergent policies. Replacing the maximum entropy regularization with KL regularization against the cloned policy, i.e.,  $\mathbb{D}_{\mathrm{KL}}[q_{\theta}(\boldsymbol {a}\mid s)\parallel q_{\theta_1}(\boldsymbol {a}\mid s)]$  , mitigated divergent behavior. This regularization still retains a maximum entropy effect if the BC policy behaves like a stationary process. This approach matches prior work on KL-regularized RL from demonstrations, e.g., [69, 61].
'Faithful' heteroscedastic regression loss. To fit a parametric pseudo-posterior to the expert dataset, we use the predictive distribution for conditional density estimation [64] using heteroscedastic regression [63]. A practical issue with heteroscedastic regression with function approximators is the incorrect modeling of data as noise [70]. This can be overcome with a 'faithful' loss function and modelling construction [71], which achieves the desired minimization of the squared error in the mean and fits the predictive variance to model the residual errors. For more details, see Appendix D.
Refining the coherent reward. The HETSTAT network we use still approximates a stationary process, as shown in Figure 3. To ensure spurious reward values from approximation errors are not exploited during learning, it can be beneficial to refine, in a minimax fashion, the coherent reward with the additional data seen during learning. This minimax refinement both reduces the stationary approximation errors of the policy, while also improving the reward from the game-theoretic IRL perspective. For further details, intuition and ablations, see Appendix E, G and N.6 respectively.
Algorithm 1 with these additions can be found summarized in Algorithm 2 in Appendix I. An extensive ablation study of these adjustments can be found in Appendix N.
5 Experimental results
We evaluate CSIL against baseline methods on tabular and continuous state-action environments. The baselines are popular entropy-regularized imitation learning methods discussed in Section 2. Moreover, ablation studies are provided in Appendix N for the experiments in Section 5.2 and 5.3
MDP	Variant	Expert	BC	Classifier	ME-IRL	GAIL	IQLearn	PPIL	CSIL
Dense	Nominal	0.266	0.200	0.249	0.251	0.253	0.244	0.229	0.257
Windy	0.123	0.086	0.103	0.111	0.105	0.104	0.104	0.107
Sparse	Nominal	1.237	1.237	1.237	1.131	1.237	1.237	1.237	1.237
Windy	0.052	0.002	0.044	0.036	0.043	0.002	0.002	0.044
Table 1: Inverse optimal control, combining SPI and known dynamics, in tabular MDPs. The 'dense' MDP has a uniform initial state distribution and four goal states. The 'sparse' MDP has one initial state, one goal state and one forbidden state. The agents are trained on the nominal MDP. The 'windy' MDP has a random disturbance across all states to evaluate the robustness of the policy. CSIL performs well across all settings despite being the simpler algorithm relative to the IRL baselines.
Figure 4: Normalized performance of CSIL against baselines for online imitation learning for MuJoCo Gym tasks. Uncertainty intervals depict quartiles over ten seeds. To assess convergence across seeds, performance is chosen using the highest 25th percentile of the episode return during learning.
5.1 Tabular inverse optimal control
We consider inverse optimal control in two tabular environments to examine the performance of CSIL without approximations. One ('dense') has a uniform initial distribution across the state space and several goal states, while the other ('sparse') has one initial state, one goal state, and one forbidden state with a large negative reward. Table 1 and Appendix M.1 shows the results for a nominal and a 'windy' environment, which has a fixed dynamics disturbance to test the policy robustness. The results show that CSIL is effective across all settings, especially the sparse environment where many baselines produce policies that are brittle to disturbances and, therefore, covariate shift. Moreover, CSIL produces value functions similar to GAIL (Figures 23 and 27) while being a simpler algorithm.
5.2 Continuous control from agent demonstrations
A standard benchmark of deep imitation learning is learning MuJoCo [72] Gym [73] and Adroit [74] tasks from agent demonstrations. We evaluate online and offline learning, where a fixed dataset is used in lieu of environment interactions [75]. Acme [76] was used to implement CSIL and baselines, and expert data was obtained using r1ds [77] from existing sources [11, 74, 75]. Returns are normalized with respect to the reported expert and random performance [11] (Table 2).
Online imitation learning. In this setting, we used DAC (actor-critic GAIL), IQLearn, PPIL, SQIL, and PWIL as entropy-regularized imitation learning baselines. We evaluate on the standard benchmark of locomotion-based Gym tasks, using the SAC expert data generated by Orsini et al. [11]. In this setting, Figures 4 and 28 show CSIL closely matches the best baseline performance across environments and dataset sizes. We also evaluate on the Adroit environments, which involve manipulation tasks with a complex 27-dimensional robot hand [74]. In this setting, Figures 5 and 29 show that saddle-point methods struggle due to the instability of the optimization in high dimensions without careful regularization and hyperparameter selection [11]. In contrast, CSIL is very effective, matching or surpassing BC, highlighting the benefit of coherency for both policy initialization and improvement. In the Appendix, Figure 29 includes SAC from demonstrations (SACfd) [78, 74] as an oracle baseline with access to the true reward. CSIL exhibits greater sample efficiency than SACfd, presumably due to the BC initialization and the shaped reward, and often matches final performance. Figures 38, 39 and 40 in the Appendix show an ablation of IQLearn and PPIL with BC pre-training. We observe a fast unlearning of the BC policy due to the randomly initialized rewards, which was also observed by Orsini et al. [11], so any initial performance improvement is negligible or temporary.
Figure 5: Normalized performance of CSIL against baselines for online imitation learning for Adroit tasks. Uncertainty intervals depict quartiles over ten seeds. To assess convergence across seeds, performance is chosen using the highest filtered 25th percentile of the return during learning.
Figure 6: Normalized performance of CSIL against baselines for offline imitation learning for Gym tasks. Uncertainty intervals depict quartiles over ten seeds. To assess convergence across seeds, performance is chosen using the highest filtered 25th percentile of the episode return during learning.
Offline imitation learning. Another challenging setting for deep RL is offline learning, where the online interactions are replaced with a static 'supplementary' dataset. Note that prior works have used 'offline' learning to describe the setting where only the demonstrations are used for learning [41, 79, 42, 43], like in BC, such that the IL method becomes a form of regularized BC [80]. Standard offline learning is challenging primarily because the critic approximation can favor unseen actions, motivating appropriate regularization (e.g., [81]). We use the full-replay datasets from the d4r1 benchmark [75] of the Gym locomotion tasks for the offline dataset, so the offline and demonstration data is not the same. We also evaluate SMODICE [82] and DemoDICE [83], two offline imitation learning methods which use state-based value functions, weighted BC policy updates, and a discriminator-based reward. We opt not to reproduce them in acme and use the original implementations. As a consequence, the demonstrations are from a different dataset, but are normalized appropriately. For further details on SMODICE and DemoDICE, see Appendix A and L.
Figures 6 and 30 demonstrate that offline learning is significantly harder than online learning, with no method solving the tasks with few demonstrations. This can be attributed to the lack of overlap between the offline and demonstration data manifesting as a sparse reward signal. However, with more demonstrations, CSIL can achieve reasonable performance. This performance is partly due to the strength of the initial BC policy, but CSIL can still demonstrate policy improvement in some cases (Figure 40). Figure 30 includes CQL [81] as an oracle baseline, which has access to the true reward function. For some environments, CSIL outperforms CQL with enough demonstrations. The DemoDICE and SMODICE baselines are both strong, especially for few demonstrations. However, performance does not increase so much with more demonstrations. Since CSIL could also be implemented with a state-based value function and a weighed BC policy update (e.g., [16]), further work is needed to determine which components matter most for effective offline imitation learning.
5.3 Continuous control from human demonstrations from states and images
As a more realistic evaluation, we consider learning robot manipulation tasks such as picking (Lift), pick-and-place (PickPlaceCan), and insertion (NutAssemblySquare) from random initializations and mixed-quality human demonstrations using the robomimic datasets [84], which also include image observations. The investigation of Mandlekar et al. [84] considered offline RL and various forms of BC with extensive model selection. Instead, we investigate the applicability of imitation learning, using online learning with CSIL as an alternative to BC model selection. One aspect of these tasks is that they have a sparse reward based on success, which can be used to define an absorbing
Figure 7: Average success rate over 50 evaluations for online imitation learning for robomimic tasks. Uncertainty intervals depict quartiles over ten seeds. To assess convergence across seeds, performance is chosen using the highest 25th percentile of the averaged success during learning.
Figure 8: Average success rate over 50 evaluations for image-based online imitation learning for robomimic tasks. Uncertainty intervals depict quartiles over five seeds. Baseline results obtained from Mandlekar et al. [84]. BC (CSIL) denotes the performance of CSIL's initial policy for comparison.
state. As observed in prior work [40], in practice there is a synergy between absorbing states and the sign of the reward, where positive rewards encourage survival and negative rewards encourage minimum-time strategies. We found that CSIL was initially ill-suited to these goal-oriented tasks as the rewards are typically positive. In the same way that the AIRL reward can be designed for a given sign [11], we can design negative CSIL rewards by deriving an upper-bound of the log policy ratio and subtracting it in the reward. For further details, see Appendix F. An ablation study in Appendix N.9 shows that the negative CSIL reward matches or outperforms a constant negative reward, especially when given fewer demonstrations. These tasks are also more difficult due to the variance in the initial state, requiring much larger function approximators than in Section 5.2.
Figures 7 and 36 show the performance of CSIL and baselines, where up to 200 demonstrations are required to sufficiently solve the tasks. CSIL achieves effective performance and reasonable demonstration sample efficiency across all environments, while baselines such as DAC struggle to solve the harder tasks. Moreover, Figure 36 shows CSIL is again more sample efficient than SACfd.
Figure 8 shows the performance on CSIL when using image-based observations. A shared CNN torso is used between the policy, reward and critic, and the convolutional layers are frozen after the BC stage. CSIL was able to solve the simpler tasks in this setting, matching the model selection strategy of Mandlekar et al. [84], demonstrating its scalability. In the offline setting, Figure 37 shows state-based results with sub-optimal ('mixed human') demonstrations as supplementary dataset. While some improvement could be made on simpler tasks, on the whole it appears much harder to learn from suboptimal human demonstrations offline. This supports previous observations that human behaviour can be significantly different to that of RL agents, such that imitation performance is affected [11, 84].
We have shown that 'coherency' is an effective approach to IL, combining BC with IRL-based finetuning by using a shaped learned reward for which the BC policy is optimal. We have demonstrated the effectiveness of CSIL empirically across a range of settings, particularly for high-dimensional tasks and offline learning, due to CSIL leveraging BC for initializing the policy and reward.