SlowGuess/NeurIPS_2023_OCR huggingface dataset:
In Figure 35 of the Appendix, we investigate why baselines IQLearn and PPIL, both similar to CSIL, struggle in the high-dimensional and offline settings. Firstly, the sensitivity of the saddle-point optimization can be observed in the stability of the expert reward and critic values during learning. Secondly, we observe that the policy optimization does not effectively minimize the BC objective by fitting the expert actions. This suggests that these methods do not necessarily converge to a BC solution and lack the 'coherence' quality that CSIL has that allows it to refine BC policies.
A current practical limitation of CSIL is understanding when the initial BC policy is viable as a coherent reward. Our empirical results show that CSIL can solve complex control tasks using only one demonstration, where its BC policy is ineffective. However, if a demonstration-rich BC policy cannot solve the task (e.g., image-based NutAssemblySquare), CSIL also appears to be unable to solve the task. This suggests there is scope to improve the BC models and learning to make performance more consistent across more complex environments. An avenue for future work is to investigate the performance of CSIL with richer policy classes beyond MLPs, such as recurrent architectures and multi-modal action distributions, to assess the implications of the richer coherent reward and better model sub-optimal human demonstrations.
We wish to thank Jost Tobias Springenberg, Markus Wulfmeier, Todor Davchev, Ksenia Konyushkova, Abbas Abdelmaleki and Matthew Hoffman for helpful discussions during the project. We would also like to thank Gabriel Dulac-Arnold and Sabela Ramos for help with setting up datasets and experiments, Luca Viano for help reproducing PPIL, and Martin Riedmiller, Markus Wulfmeier, Oleg Arenz, Davide Tateo, Boris Belousov, Michael Lutter and Gokul Swarmy for feedback on previous drafts. We also thank the wider Google DeepMind research and engineering teams for the technical and intellectual infrastructure upon which this work is built, in particular the developers of acme. Some baseline experiments were run on the computer cluster of the Intelligence Autonomous Systems group at TU Darmstadt, which is maintained by Daniel Palenicek and Tim Schneider. Joe Watson acknowledges the grant "Einrichtung eines Labors des Deutschen Forschungszentrum für Künstliche Intelligenz (DFKI) an der Technischen Universität Darmstadt" of the Hessian Ministry of Science and Art.
[1] Takayuki Osa, Joni Pajarinen, Gerhard Neumann, J Andrew Bagnell, Pieter Abbeel, Jan Peters, et al. An algorithmic perspective on imitation learning. Foundations and Trends® in Robotics, 2018.
[2] Stefan Schaal. Is imitation learning the route to humanoid robots? Trends in Cognitive Sciences, 1999.
[3] D. A. Pomerleau. Efficient training of artificial neural networks for autonomous navigation. Neural Computation, 1991.
[4] Rudolph E. Kalman. When is a linear control system optimal? Journal of Basic Engineering, 1964.
[5] Andrew Y. Ng and Stuart J. Russell. Algorithms for inverse reinforcement learning. In International Conference on Machine Learning, 2000.
[6] P. Abbeel and A. Y. Ng. Apprenticeship learning via inverse reinforcement learning. In International Conference on Machine Learning, 2004.
[7] Martin L Puterman. Markov Decision Processes: Discrete Stochastic Dynamic Programming. John Wiley & Sons, 2014.
[8] Richard S Sutton and Andrew G Barto. Reinforcement learning: An introduction. MIT press, 2018.
[9] Stéphane Ross and Drew Bagnell. A reduction of imitation learning and structured prediction to no-regret online learning. In International Conference on Artificial Intelligence and Statistics, 2011.
[10] Boyuan Zheng, Sunny Verma, Jianlong Zhou, Ivor W Tsang, and Fang Chen. Imitation learning: Progress, taxonomies and challenges. IEEE Transactions on Neural Networks and Learning Systems, 2022.
[11] Manu Orsini, Anton Raichuk, Léonard Hussenot, Damien Vincent, Robert Dadashi, Sertan Girgin, Matthieu Geist, Olivier Bachem, Olivier Pietquin, and Marcin Andrychowicz. What matters for adversarial imitation learning? In Advances in Neural Information Processing Systems, 2021.
[12] Rohit Jena, Changliu Liu, and Katia Sycara. Augmenting GAIL with BC for sample efficient imitation learning. In Conference on Robot Learning, 2020.
[13] Jan Peters, Katharina Mulling, and Yasemin Altun. Relative entropy policy search. In Proceedings of the AAAI Conference on Artificial Intelligence, 2010.
[14] Roy Fox, Ari Pakman, and Naftali Tishby. Taming the noise in reinforcement learning via soft updates. In Conference on Uncertainty in Artificial Intelligence, 2016.
A Extended related work
This section discusses the relevant prior literature in more breadth and depth.
Generative adversarial imitation learning (GAIL) [37] combines the game-theoretic IRL objective with causal entropy regularization and convex regularization of the reward. This primal optimization problem can be transformed into a dual objective, which has the interpretation of minimizing a divergence between the expert and imitator's stationary distributions, where the divergence corresponds to the choice of convex reward regularization. Choosing regularization that reduced the reward to a classification problem corresponds to the Jensen-Shannon (JS) divergence. While GAIL used on-policy RL, DAC [40] used off-policy actor-critic methods (e.g., SAC) in an ad hoc fashion for greater sample efficiency. ORIL [85] adopts a classification-based reward in the offline setting. Kostrikov et al. [40] also note that the GAIL classifier can be used in different forms to produce rewards with different forms. AIRL [65] highlights that the GAIL reward is shaped by the dynamics if the MDP following the shaping theory of Ng and Russell [5], and use additional approximators that disentangle the shaping terms. NAIL [79] replaces the minimax nature of adversarial learning with a more stable max-max approach using expectation maximization-like optimization on a lower-bound of the KL divergence between policies. Its reward function is defined, as with AIRL, as the log density ratio of the stationary distributions of the expert and policy, regularized by the causal entropy. Therefore, NAIL and CSIL have similar reward structures; however, while CSIL's reward can be obtained from BC and prior design, NAIL requires iterative density ratio estimation during learning. Ghasemipour et al. [39] investigate the performance of AIRL with different  $f$ -divergences and corresponding reward regularization beyond JS such as the forward and reverse KL divergence. Ghasemipour et al. [39] discuss the connection between the log density ratio, AIRL, GAIL and their proposed methods. Ghasemipour et al. [39] also connect BC to IRL, but from the divergence perspective rather than the likelihood perspective, as BC can be cast as the conditional KL divergence between the expert and agent policies.
A separate line of study uses the convex duality in a slightly different fashion for 'discounted distribution correction estimation' (DICE) methods. ValueDICE uses the Donsker-Varadhan representation to transform the KL minimization problem into an alternative form and a change of variables to represent the log ratio of the stationary distributions as with an implicit representation using the Bellman operator. However, this alternative form is still a saddle-point problem like GAIL, but the final objective is more amenable to off-policy optimization. DemoDICE [83] considers the offline setting and incorporates KL-regularization into the DICE framework, simplifying optimization. Due to the offline setting, the reward is obtained using a pre-trained classifier. Policy improvement reduces to weighted BC, like in many KL-regularized RL methods [86, 16]. SmoDICE [82] is closely related to DemoDICE, but focuses on matching the state distribution and, therefore, does not rely on expert actions, so the discriminator-based reward depends only on the state.
Inverse soft  $Q$ -learning (IQLearn) [42] combines the game-theoretic IRL objective, entropy-regularized actor-critic (SAC), convex reward regularization and implicit rewards to perform IRL using only the critic by maximizing the expert's implicit reward. While Garg et al. [42] claim their method is non-adversarial because it does not use a classifier, it is still derived from the game-theoretic objective. For continuous control, the IQLearn implementation requires shaping the implicit reward using online samples, making learning not so different to GAIL-based algorithms. Moreover, the implementation replaces the minimization of the initial value with the value evaluated for the expert samples. For SAC, in which the value function is computed by evaluating the critic with the current policy, this shapes the critic to have a local maximum around the expert samples. We found this implementation detail crucial to reproduce the results for a few demonstrations and discuss further in Appendix H.
Proximal point imitation learning (PPIL) [43] combines the game-theoretic and linear programming forms of IRL and entropy-regularized actor-critic RL to derive an effective method with convergence guarantees. Its construction yields a single objective (the dual) for optimizing the reward and critic jointly in a more stable fashion, rather than alternating updates like in DAC, and unlike IQLearn, it has an explicit reward model. For continuous control, its implementation builds on SAC and IQLearn, including the aforementioned tricks, but with two key differences: an explicit reward function and minimization of the 'logistic', rather than squared, Bellman error [19].
Compared to the works discussed above, CSIL has two main distinctions. One is the reward hypothesis space in the form of the log policy ratio, derived from KL-regularized RL, as opposed to a classifier-
based reward or critic-as-classifier implicit reward. Secondly, the consequence of this coherent reward is that it can be pre-trained using BC, and the BC policy can be used as a reference policy. While BC pre-training could be incorporated into prior methods in an ad-hoc fashion, their optimization objectives do not guarantee that coherency is maintained. This is demonstrated in Figure 35, where IQLearn and PPIL combined with BC pre-training and KL regularization do not converge like CSIL.
Concurrent work, least-squares inverse  $Q$ -learning (LS-IQ) [87], addresses the undocumented regularization of IQLearn by minimizing a mixture divergence, which yields the desired regularization terms and better treatment of absorbing states. The treatment moves away from the entropy-regularized RL view and instead opts for explicit regularization critics.
Szot et al. [88], in concurrent work, use the bi-level optimization view of IRL to optimize the maximum entropy reward through an upper-level behavioral cloning loss using 'meta' gradients. This approach achieves coherency through the top-level loss while regularizing the learned reward. However, this approach is restricted to on-policy RL algorithms that directly use the reward rather than actor-critics methods that use a  $Q$  function, making it less sample efficient.
Brantley et al. [24] propose disagreement-regularized imitation learning, which trains an ensemble of policies via BC and uses their predictive variance to define a proxy reward. This work shares the same motivation as CSIL in achieving coherency and tackling the covariate shift problem. Moreover, using ensembles is similar to CSIL's stationary policies, as both models are inspired by Bayesian modeling and provide uncertainty quantification. Unlike CSIL, the method is not motivated by reward shaping, and the reward is obtained by thresholding the predictive variance to  $\pm 1$  rather than assume that the likelihood values are meaningful.
Taranovic et al. [89] derive an AIRL formulation from the KL divergence between expert and policy that results in a reward consisting of log policy ratios. However, they use a classifier to estimate these quantities and ignore some policy terms in the implementation.
Swamy et al. [90] use the BC policy for replay estimation [25] and a membership classifier to augment the expert dataset used during IRL.
'Noisy' BC [91] is tangentially related to CSIL, as it uses the BC log-likelihood as a proxy reward to finetune the BC policy using policy gradient updates when learning from demonstrations that are sub-optimal due to additive noise.
Soft  $Q$ -imitation learning (SQIL) [23] uses binary rewards for the demonstration and interaction transitions, bypassing the complexity of standard IRL. The benefit of this approach comes from the critic learning to overcome the covariate shift problem through the credit assignment. This approach has no theoretical guarantees, and the simplistic nature of the reward does not encourage stable convergence in practice.
Barde et al. [51] proposed adversarial soft advantage fitting (ASAF), an adversarial approach that constructs the classifier using two policies over trajectories or transitions. This approach is attractive as the classification step performs the policy update, so no policy evaluation step is required. While this approach has clear connections to regularized behavioral cloning, it's uncertain how the policy learns to overcome covariate shift without any credit assignment. Moreover, this method has no convergence guarantees.
CSIL adopts a Bayesian-style approach to its initial BC policy by specifying a prior policy but uses the predictive distribution rather than a likelihood to fit the demonstration data. This is in contrast to Bayesian IRL [92], where the prior is placed over the reward function as a means of reasoning about its ambiguity. As the reward is an unobserved abstract quantity, the exponentiated critic is adopted as a pseudo-likelihood, and approximate inference is required to estimate the posterior reward. Using a pseudo-likelihood means that Bayesian IRL and ME-IRL are not so different in practice, especially when using point estimates of the posterior [93].
B Likelihoods and priors in imitation learning
The pseudo-posterior perspective in this work was inspired by van Campenhout and Cover [56], who identify that effective likelihoods arise from relevant functions that describe a distribution. With this in mind, we believe it is useful to see IRL as extending BC from a myopic, state-independent likelihood to one that encodes the causal, state-dependent nature of the problem as done ME-IRL, essentially incorporating the structure of the MDP into the regression problem [50]. This perspective has two key consequences. One is the importance of the prior, which provides important regularization. For example, Ziniu et al. [80] have previously shown that the ValueDICE does not necessarily outperform BC as previously reported but matches BC with appropriate regularization. Moreover, the MIMIC-EXP formulation of BC of Rajaraman et al. [25] proposes a BC policy with a stationary prior,
$$
\pi (\boldsymbol {a} \mid \boldsymbol {s}) = \left\{ \begin{array}{l} p _ {\mathcal {D}} (\boldsymbol {a} \mid \boldsymbol {s}) \text {i f} \boldsymbol {s} \in \mathcal {D}, \text {w h e r e} p _ {\mathcal {D}} \text {d e n o t e s c o n d i t i o n a l d e n s i t y e s t i m a t i o n o f} \mathcal {D}, \\ \mathcal {U} _ {\mathcal {A}} (\boldsymbol {a}) \text {o t h e r w i s e}. \end{array} \right.
$$
This work shows how policies like these can be implemented for continuous control using stationary architectures.
The second consequence is the open question of how to define effective likelihoods for imitation learning. ME-IRL and other prior works use feature matching, which, while an appropriate constraint, raises the additional question of which features are needed. Using a classifier in AIRL has proved highly effective, but it is not always straightforward to implement in practice without regularization and hyperparameter tuning. Much of IQLearn's empirical success for continuous tasks arises from shaping the likelihood via the critic through the expert and on-policy regularization terms.
This work addresses the likelihood design problem through coherence and the inversion analysis in Theorem 1, which informs the function-approximated critic regularization described in Section 4. An open question is how these design choices could be further improved, as the ablation study in Figure 46 shows that the benefit of the critic regularization is not evident for all tasks.
C Parametric stationary processes
The effectiveness of CSIL is down to its stationary policy design. This section briefly summarizes the definition of stationary processes and how they can be approximated with MLPs following the work of Meronen et al. [68].
While typically discussed for temporal stochastic processes whose input is time, stationary processes (Definition 3) are those whose statistics do not vary with their input. In machine learning, these processes are useful as they allow us to design priors that are consistent per input but capture correlations (e.g., smoothness) across inputs, which can be used for regularization in regression and classification [62]. In the continuous setting, stationarity can be defined for linear models, i.e.,  $y = \boldsymbol{w}^{\top}\phi(\boldsymbol{x})$ ,  $\boldsymbol{w} \sim \mathcal{N}(\boldsymbol{\mu}, \boldsymbol{\Sigma})$ , through their kernel function  $\mathcal{K}(\boldsymbol{x},\boldsymbol{x}^{\prime}) = \phi(\boldsymbol{x})^{\top}\phi(\boldsymbol{x}^{\prime})$ , the feature inner product. For stationary kernels, the kernel depends only on a relative shift  $r \in \mathbb{R}^{d}$  between inputs, i.e.  $\mathcal{K}(r) = \phi(x)^{\top}\phi(x + r)$ . By combining the kernel with the weight prior of the linear model, we can define the covariance function, e.g.  $\mathcal{C}(r) = \phi(\boldsymbol{x})^{\top}\Sigma\phi(\boldsymbol{x} + r)$ . Boschner's theorem [94] states that the covariance function of a stationary process can be represented as the Fourier transform of a positive finite measure. If the measure has a density, it is known as the spectral density  $S(\omega)$ , and it is a Fourier dual of the covariance function if the necessary conditions apply, known as the Wiener-Khinchin theorem [95],
$$
\mathcal {C} (\boldsymbol {r}) = \frac {1}{(2 \pi) ^ {d}} \int_ {\mathbb {R} ^ {d}} S (\boldsymbol {\omega}) \exp (i \boldsymbol {\omega} ^ {\top} \boldsymbol {r}) \mathrm {d} \boldsymbol {r}, \quad S (\boldsymbol {\omega}) = \int_ {\mathbb {R} ^ {d}} \mathcal {C} (\boldsymbol {r}) \exp (- i \boldsymbol {\omega} ^ {\top} \boldsymbol {r}) \mathrm {d} \boldsymbol {\omega}. \tag {9}
$$
To use this theory to design stationary function approximators, we can use the Wiener-Khintchin theorem to design function approximators that produce a Monte Carlo approximation of stationary kernels. Firstly, we consider only the last layer of the MLP, such that it can be viewed as a linear model. Secondly, Meronen et al. [68] describe several periodic activation functions that produce products of the form  $\exp (i\omega^{\top}r)$ , such that the inner product of the whole feature space performs a Monte Carlo approximation of the integral. We implemented sinusoidal, triangular, and periodic ReLU activations (see Figure 3 of Meronen et al. [68]) as a policy architecture hyperparameter. Thirdly, the input into the stationary kernel approximation should be small (e.g.,  $< 20$ ), as the stationary behavior
Heteroscedastic MLP (NLLH) Train MSE  $= 0.11$
S
S
Figure 9: We repeat the regression problem from Figure 3 to highlight the effect of 'faithful' heteroscedastic regression. The NLLH objective has a tendency to model 'signal' as noise, which is most clearly seen in the heteroscedastic MLP. The faithful heteroscedastic loss reduces the mean prediction error significantly without affecting uncertainty quantification for both models.
is dictated by  $\boldsymbol{r}$ , and this distance becomes less meaningful in higher dimensions. This means the MLP feature space that feeds into the last layer should be small or have a 'bottleneck' architecture to compress the internal representation. Finally, we require a density to represent the spectral density, which also defines the nature of the stationary process. We use a Gaussian, but other distributions such as a Student-  $t$  or Cauchy could be used if 'rougher' processes are desired. In summary, the last layer activations take the form  $f_{p}(\boldsymbol{W}\phi(\boldsymbol{x}))$ , where  $f_{p}$  is the periodic activation function of choice,  $\boldsymbol{W}$  are weights sampled from the distribution of choice and  $\phi(\boldsymbol{x})$  is an arbitrary (but bottlenecked) MLP feature space of choice.
Related work combining stationary kernel approximations with MLPs have constrained the feature space with a residual architecture and spectral normalization to bound the upper and lower Lipschitz constant [96], to prevent these features overfitting. We found we did not need this to achieve approximate stationarity (e.g., Figure 3), which avoided the need for input-dependent policy architectures which may lead to a risk of underfitting.
A crucial practical detail is that the sampled weights  $W$  should be trained and not kept constant, as the theory suggests. In practice, we did not observe the statistics of these weights changing significantly during training, while training this weight layer reduced underfitting enough for substantially more effective BC performance.
D Faithful heteroscedastic regression
This section summarizes the approach of Stirn et al. [71]. The faithful objective combines the mean squared error (MSE) and negative log-likelihood NLLH loss with a careful construction of the heteroscedastic model. The loss is
$$
\mathcal {L} (\mathcal {D}, \boldsymbol {\theta}) = \mathbb {E} _ {\boldsymbol {s}, \boldsymbol {a} \sim \mathcal {D}} \left[ \left(\boldsymbol {a} - \boldsymbol {\mu} _ {\boldsymbol {\theta}} (\boldsymbol {s})\right) ^ {2} - \log \tilde {q} (\boldsymbol {a} \mid \boldsymbol {s}) \right], \quad \tilde {q} (\boldsymbol {a} \mid \boldsymbol {s}) = \mathcal {N} (\operatorname {s g} (\boldsymbol {\mu} _ {\boldsymbol {\theta}} (\boldsymbol {s})), \boldsymbol {\Sigma} _ {\boldsymbol {\theta}} (\boldsymbol {s})), \tag {10}
$$
where  $\mathfrak{sg}(\cdot)$  denotes the stop gradient operator. Moreover, for the models 'shared torso' (e.g., features), the gradient is stopped between the torso and the predictive variance so that the features are only trained to satisfy the MSE objective. Figure 9 illustrates the impact of the alternative loss function and model adjustments. Note that the faithful objective is not strictly necessary for BC or CSIL and is mitigating potential under-fitting previously observed empirically. The issue of modeling 'signal' as noise when using heteroscedastic function approximators is poorly understood and depends on the model architecture, initialization, and the dataset considered. While motivated by 1D toy examples, it was observed during development that the faithful objective is beneficial for CSIL's downstream performance.
E Reward refinement and function-space relative entropy estimation
In the function approximation setting, we refine the reward model with additional non-expert samples. Motivated by Theorem 2, we maximize the objective using samples from the replay buffers
$$
\mathcal {J} _ {r} (\boldsymbol {\theta}) = \mathbb {E} _ {\boldsymbol {s}, \boldsymbol {a} \sim \mathcal {D}} \left[ \alpha \log \frac {q _ {\boldsymbol {\theta}} (\boldsymbol {a} \mid \boldsymbol {s})}{p (\boldsymbol {a} \mid \boldsymbol {s})} \right] - \mathbb {E} _ {\boldsymbol {s}, \boldsymbol {a} \sim \rho_ {\pi}} \left[ \alpha \log \frac {q _ {\boldsymbol {\theta}} (\boldsymbol {a} \mid \boldsymbol {s})}{p (\boldsymbol {a} \mid \boldsymbol {s})} \right]. \tag {11}
$$