Title: Coherent Soft Imitation Learning
Venue: NeurIPS 2023
Accepted: yes
Keywords: imitation learning

inverse reinforcement learning

behavioral cloning

learning from demonstration
Abstract: Imitation learning methods seek to learn from an expert either through behavioral cloning (BC) for the policy or inverse reinforcement learning (IRL) for the reward.
Such methods enable agents to learn complex tasks from humans that are difficult to capture with hand-designed reward functions.
Choosing between BC or IRL for imitation depends on the quality and state-action coverage of the demonstrations, as well as additional access to the Markov decision process. 
Hybrid strategies that combine BC and IRL are rare, as initial policy optimization against inaccurate rewards diminishes the benefit of pretraining the policy with BC.
Our work derives an imitation method that captures the strengths of both BC and IRL.
In the entropy-regularized (`soft') reinforcement learning setting, we show that the behavioral-cloned policy can be used as both a shaped reward and a critic hypothesis space by inverting the regularized policy update. 
This coherency facilitates fine-tuning cloned policies using the reward estimate and additional interactions with the environment.
This approach conveniently achieves imitation learning through initial behavioral cloning and subsequent refinement via RL with online or offline data sources.
The simplicity of the approach enables graceful scaling to high-dimensional and vision-based tasks, with stable learning and minimal hyperparameter tuning, in contrast to adversarial approaches.
For the open-source implementation and simulation results, see https://joemwatson.github.io/csil/.

Coherent Soft Imitation Learning

Abstract

Imitation learning methods seek to learn from an expert either through behavioral cloning (BC) for the policy or inverse reinforcement learning (IRL) for the reward. Such methods enable agents to learn complex tasks from humans that are difficult to capture with hand-designed reward functions. Choosing between BC or IRL for imitation depends on the quality and state-action coverage of the demonstrations, as well as additional access to the Markov decision process. Hybrid strategies that combine BC and IRL are rare, as initial policy optimization against inaccurate rewards diminishes the benefit of pretraining the policy with BC. This work derives an imitation method that captures the strengths of both BC and IRL. In the entropy-regularized ('soft') reinforcement learning setting, we show that the behavioral-cloned policy can be used as both a shaped reward and a critic hypothesis space by inverting the regularized policy update. This coherency facilitates fine-tuning cloned policies using the reward estimate and additional interactions with the environment. Our approach conveniently achieves imitation learning through initial behavioral cloning and subsequent refinement via RL with online or offline data sources. The simplicity of the approach enables graceful scaling to high-dimensional and vision-based tasks, with stable learning and minimal hyperparameter tuning, in contrast to adversarial approaches. For the open-source implementation and simulation results, see joemwatson.github.io/csil.

s

a

Reward

s

Classifier

s

Ours

1 Introduction

Imitation learning (IL) methods [1] provide algorithms for learning skills from expert demonstrations in domains such as robotics [2], either through direct mimicry (behavioral cloning, BC [3]) or inferring the latent reward (inverse reinforcement learning, IRL [4-6]) assuming a Markov decision process (MDP) [7, 8]. While BC is straightforward to implement, it is susceptible to covariate shift during evaluation [9]. IRL overcomes this problem but tackles the more complex task of jointly learning

dfki.de

a reward and policy from interactions. So far, these two approaches to imitation have been largely separated [10], as optimizing the policy with an evolving reward estimate counteracts the benefit of pre-training the policy [11]. Combining the objectives would require careful weighting during learning and lacks convergence guarantees [12]. The benefit of combining BC and IRL is sampleefficient learning that enables any BC policy to be further improved using additional experience.

Contribution. Our IL method naturally combines BC and IRL by using the perspectives of entropyregularized reinforcement learning (RL) [13-19] and gradient-based IRL [20, 21]. Our approach uses a BC policy to define an estimate of a shaped reward for which it is optimal, which can then be used to finetune the policy with additional knowledge, such as online interactions, offline data, or a dynamics model. Using the cloned policy to specify the reward avoids the need for careful regularization and hyperparameter tuning associated with adversarial imitation learning [11, 22], and we estimate a shaped variant of the true reward, rather than use a heuristic proxy, e.g., [23, 24]. In summary,

We use entropy-regularized RL to obtain a reward for which a behavioral-cloned policy is optimal, which we use to improve the policy and overcome the covariate shift problem.

We introduce approximate stationary stochastic process policies to implement this approach for continuous control by constructing specialized neural network architectures.

We show strong performance on online and offline imitation for high-dimensional and image-based continuous control tasks compared to current state-of-the-art methods.

2 Background & related work

A Markov decision process is a tuple ⟨S , A , P , r, γ, µ 0 ⟩ , where S is the state space, A is the action space, P : S ×A×S → R + is the transition model, r : S ×A→ R is the reward function, γ is the discount factor, and µ 0 is the initial state distribution s 0 ∼ µ 0 ( · ) . When evaluating a policy π ∈ Π , we use the occupancy measure ρ π ( s , a ) ≜ π ( a | s ) µ π ( s ) , µ π ( s ) ≜ ∑ ∞ t =0 γ t µ π t ( s ) where µ π t +1 ( s ′ ) = ∫ P ( s ′ | s , a ) π ( a | s ) µ π t ( s ) d s d a . This measure is used to compute infinitehorizon discounted expectations, E s , a ∼ ρ π [ f ( s , a )] = E s t +1 ∼P ( ·| s t , a t ) , a t ∼ π ( ·| s t ) [ ∑ ∞ t =0 γ t f ( s t , a t )] . Densities d π ( s , a ) = (1 -γ ) ρ π ( s , a ) and ν π ( s ) = (1 -γ ) µ π ( s ) are normalized measures. When the reward is unknown, imitation can be performed from a dataset D of transitions ( s , a , s ′ ) as demonstrations, obtained from the discounted occupancy measure of the expert policy π ∗ = arg max π E s , a ∼ ρ π [ r ( s , a )] . The policy could be inferred directly or by inferring a reward and policy jointly, referred to as behavioral cloning and inverse reinforcement learning respectively.

Behavioral cloning. The simplest approach to imitation is to directly mimic the behavior using regression [3], min π ∈ Π l ( π, D ) + λ Ψ( π ) , with some loss function l , hypothesis space Π and regularizer Ψ . This method is effective with sufficient data coverage and an appropriate hypothesis space but suffers from compounding errors and cannot improve unless querying the expert [9]. The MIMIC-EXP algorithm of Rajaraman et al. [25] shows that the optimal no-interaction policy matches BC.

Inverse reinforcement learning. Rather than just inferring the policy, IRL infers the reward and policy jointly and relies on access to the underlying MDP [26, 5]. IRL iteratively finds a reward for which the expert policy is optimal compared to the set of possible policies while also finding the policy that maximizes this reward function. The learned policy seeks to match or even improve upon the expert using additional data beyond the demonstrations, such as environment interactions [6]. To avoid repeatedly solving the inner RL problem, one can consider the game-theoretic approach [27], where the optimization problem is a two-player zero-sum game where the policy and reward converge to the saddle point minimizing the 'apprenticeship error' min π ∈ Π max r ∈R E D [ r ] -E ρ π [ r ] . This approach is attractive as it learns the policy and reward concurrently [28]. However, in practice, saddle-point optimization can be challenging when scaling to high dimensions and can exhibit instabilities and hyperparameter sensitivity [11]. Compared to BC, IRL methods in theory scale to longer task horizons by overcoming compounding errors through environment interactions [29].

Maximum entropy inverse reinforcement learning & soft policy iteration. IRL is an underspecified problem [26], in particular in regions outside of the demonstration distribution. An effective way to address this is causal entropy regularization [30], by applying the principle of maximum entropy [31] to IRL (ME-IRL). Moreover, the entropy-regularized formulation has an elegant closedform solution [30]. The approach can be expressed as an equivalent constrained minimum relative entropy problem for policy q ( a | s ) against a uniform prior p ( a | s ) = U A ( a ) using the Kullback-

Leibler (KL) divergence and a constraint matching expected features ϕ , with Lagrangian

The constraint term can be interpreted as the apprenticeship error, where w ⊤ ϕ ( s , a ) = r w ( s , a ) . Solving Equation 1 using dynamic programming yields the 'soft' Bellman equation [30],

for temperature α = 1 . Using Jensen's inequality and importance sampling, this target is typically replaced with a lower-bound [32], which has the same optimum and samples from the optimized policy rather than the initial policy like many practical deep RL algorithms [17]:

The policy update blends the exponentiated advantage function 'pseudo-likelihood' with the prior, as a form of regularized Boltzmann policy [17] that resembles a Bayes posterior [33],

These regularized updates can also be used for RL, where it is the solution to a KL-regularized RL objective, max q E s , a ∼ ρ q [ Q ( s , a )] -α D KL [ q ( a | s ) || p ( a | s )] , where the temperature α now controls the strength of the regularization. This regularized policy update is known as soft- [14, 17] or posterior policy iteration [34, 35] (SPI, PPI), as it resembles a Bayesian update. In the function approximation setting, the update is performed in a variational fashion by minimizing the reverse KL divergence between the parametric policy q θ and the critic-derived update at sampled states s [17],

Gradient-based inverse reinforcement learning. An alternative IRL strategy is a gradient-based approach (GIRL) that avoids saddle-point optimization by learning a reward function such that the BC policy's policy gradient is zero [20], which satisfies first-order optimality. However, this approach does not remedy the ill-posed nature of IRL. Moreover, the Hessian is required for the sufficient condition of optimality [36], which is undesirable for policies with many parameters.

The above objective J π can be maximized using reparameterized gradients and minibatches from replay buffer B [17]. A complete derivation of SPI and ME-IRL is provided in Appendix K.1.

Related work. Prior state-of-the-art methods have combined the game-theoretic IRL objective [27] with entropy-regularized RL. These methods can be viewed as minimizing a divergence between the expert and policy, and include GAIL [37], AIRL [38], f -MAX [39], DAC [40], ValueDICE [41], IQLearn [42] and proximal point imitation learning (PPIL, [43]). These methods differ through their choice of on-policy policy gradient (GAIL) or off-policy actor-critic (DAC, IQLearn, PPIL), and also how the minimax optimization is implemented, e.g., using a classifier (GAIL, DAC, AIRL), implicit reward functions (ValueDICE, IQLearn) or Lagrangian dual objective (PPIL). Alternative approaches to entropy-regularized IL use the Wasserstein metric (PWIL) [44], labelling of sparse proxy rewards (SQIL) [23], feature matching [30, 45-47], maximum likelihood [48] and matching state marginals [49] to specify the reward. Prior works at the intersection of IRL and BC include policy matching [50], policy-based GAIL classifiers [51], annealing between a BC and GAIL policy [12] and discriminatorweighted BC [52]. Our approach is inspired by gradient-based IRL [20, 36, 21], which avoids the game-theoretic objective and instead estimates the reward by analyzing the policy update steps with respect to the BC policy. An entropy-regularized GIRL setting was investigated in the context of learning from policy updates [53]. For further discussion on related work, see Appendix A.

3 Coherent imitation learning

For efficient and effective imitation learning, we would like to combine the simplicity of behavioral cloning with the structure of the MDP, as this would provide a means to initialize and refine the imitating policy through interactions with the environment. In this section, we propose such a method using a GIRL-inspired connection between BC and IRL in the entropy-regularized RL setting. By inverting the entropy-regularized policy update in Equation 4, we derive a reward for which the behavioral-cloning policy is optimal, a quality we refer to as coherence .

Inverting soft policy iteration and reward shaping. The α -regularized closed-form policy based on Equation 4 can be inverted to provide expressions for the critic and reward (Theorem 1).

Theorem 1. (KL-regularized policy improvement inversion). Let p and q α be the prior and pseudoposterior policy given by posterior policy iteration (Equation 4). The critic can be expressed as

Substituting into the KL-regularized Bellman equation lower-bound from Equation 2,

The V α ( s ) term is the 'soft' value function. We assume q α ( a | s ) = 0 whenever p ( a | s ) = 0 .

For the proof, see Appendix K.2. Jacq et al. [53] and Cao et al. [54] derived a similar inversion but for a maximum entropy policy iteration between two policies, rather than PPI between the prior and posterior. These expressions for the reward and critic are useful as they can be viewed as a valid form of shaping (Lemma 1) [54, 53], where α q a s p a s

Lemma 1. (Reward shaping, Ng et al. [55]). For a reward function r : S × A → R with optimal policy π ∈ Π , for a bounded state-based potential function Ψ : S → R , a reward function ˜ r shaped by the potential function satisfying r ( s , a ) = ˜ r ( s , a ) + Ψ( s ) -γ E s ′ ∼P ( ·| s , a ) [Ψ( s ′ )] has a shaped critic ˜ Q satisfying Q ( s , a ) = ˜ Q ( s , a ) + Ψ( s ) and the same optimal policy as the original reward.

(log ( | ) -log ( | )) is a shaped 'coherent' reward.

Definition 1. The shaped 'coherent' reward and critic are derived from the log policy ratio by combining Lemma 1 and Theorem 1, with value function V α ( s ) as the potential Ψ( s ) . When policy q α ( a | s ) models the data D while matching its prior p otherwise, the density ratio should exhibit the following shaping

s

a

-

0

.

3

0

.

0

0

.

3

0

.

6

0

.

9

In a continuous setting, the learned policy should capture this shaping approximately (Figure 2).

Coherent soft imitation learning (CSIL, Algorithm 1) uses the BC policy to initialize the coherent reward and uses this reward to improve the policy further with additional interactions outside of D . However, for the coherent reward to match Definition 1, the policy q ( a | s ) needs to match the policy prior p ( a | s ) outside of the data distribution. To achieve this, it is useful to recognize the policy q ( a | s ) as a 'pseudo-posterior' of prior p ( a | s ) and incorporate this Bayesian view into the BC step.

Imitation learning with pseudo-posteriors. The 'pseudo-posteriors' formalism [35, 56-59] is a generalized way of viewing the policy derived for ME-IRL and KL-regularized RL in Equation 4.

Definition 2. Pseudo-posteriors are solutions to KL -constrained or minimum divergence problems with an additional scalar objective or vector-valued constraint term and Lagrange multipliers λ ,

The exp( λ ⊤ f ( x )) term is an unnormalized 'pseudo-likelihood', as it facilities Bayesian inference to solve a regularized optimization problem specified by f .

Optimizing the policy distribution, the top objective captures KL-regularized RL, where f is the return, and the regularization is implemented as a constraint with bound ϵ or soft penalty with constant λ . The bottom objective is the form seen in ME-IRL (Equation 1), where f is the feature space that defines the reward model. Contrasting these pseudo-posterior policies against BC with Gaussian processes, e.g., [60, 61], we can compare the Bayesian likelihood used for regression with the critic-based pseudo-likelihood obtained from imitation learning. The pseudo-likelihoods in entropy-regularized IL methods produce an effective imitating policy by incorporating the MDP and trajectory distribution because f captures the cumulative reward, compared to just the action prediction error typically captured by BC regression likelihoods. This point is expanded in Appendix B.

Algorithm 1: Coherent soft imitation learning (CSIL)

Behavior cloning with pseudo-posteriors. Viewing the policy q ( a | s ) as a pseudo-posterior has three main implications when performing behavior cloning and designing the policy and prior:

We perform conditional density estimation to maximize the posterior predictive likelihood, rather than using a data likelihood, which would require assuming additive noise [62, 33].

We require a hypothesis space p ( a | s , w ) , e.g., a tabular policy or function approximator, that is used to define both the prior and posterior through weights w .

We need a fixed weight prior p ( w ) that results in a predictive distribution that matches the desired policy prior p ( a | s ) = ∫ p ( a | s , w ) p ( w ) d w ∀ s ∈ S .

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 ( a | s ) = U A ( 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],

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,

Perform regularized BC on the demonstrations with a parametric stochastic policy q θ ( a | s ) .

With temperature β < α , perform SPI to improve on the BC policy with the coherent reward.

Define the coherent reward, r θ ( s , a ) = α (log q θ ( a | s ) -log p ( a | s )) , with temperature α .

The reduced temperature β is required to improve the policy after inverting it with temperature α .

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.

s

a

Desired

s

Heteroscedastic MLP

s

Stationary Heteroscedastic MLP

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 θ ( a | s ) and coherent reward parameterization r θ ( s , a ) = α (log q θ ( a | s ) -log p ( a | s )) where α ≥ 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 ( a | s , w ) and finite parameters w ,

where λ = ( α -β ) /α and E s , a ∼D [log p ( a | s )] is constant. The regression objective bounds the IRL objective for a worst-case on-policy state distribution µ q ( s ) , which motivates its scaling through β . If D has sufficient coverage, no RL finetuning or KL regularization is required, so β = α and λ = 0 . If D does not have sufficient coverage, then let β < α so λ > 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 ( a | s ) = U A ( a ) ∀ s ∈ S , where U A is the uniform distribution over 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 : X → Y is stationary if its joint distribution in y ∈ Y , p ( y 1 , . . . , y n ) , is constant w.r.t. x 1 , . . . , x n ∈ X and all n ∈ 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 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 a = tanh ( z ( s )) , where z ( s ) = Wϕ ( s ) . The weights are factorized row-wise W = [ w 1 , . . . , w d a ] ⊤ , w i = N ( µ i , Σ i ) to define a GP with independent actions. Using change-of-variables like SAC [17], the policy is expressed per-action as

˜ W are weights drawn from a distribution, e.g., a Gaussian, that also characterizes the stochastic process and ϕ 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 ϕ E s , a ∼D [ ∇ a Q ϕ ( 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., D KL [ q θ ( a | s ) || q θ 1 ( a | 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

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.

1

3

10

30

Demonstrations

0.0

0.5

1.0

Return

1

3

10

30

Demonstrations

HalfCheetah-v2

1

3

10

30

Demonstrations

Hopper-v2

1

3

10

30

Demonstrations

bc

sqil

dac

iqlearn

ppil

pwil

csil

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 rlds [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.

1

3

10

30

Demonstrations

0.0

0.5

1.0

Return

door-v0

1

3

10

30

Demonstrations

hammer-v0

bc

sqil

dac

iqlearn

ppil

pwil

csil

1

3

10

30

Demonstrations

0.0

0.5

1.0

Return

1

3

10

30

Demonstrations

HalfCheetah-v2

1

3

10

30

Demonstrations

Hopper-v2

1

3

10

30

Demonstrations

bc (demos)

iqlearn

ppil

demodice

smodice

csil

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 d4rl 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

25

50

100

200

Demonstrations

0.0

0.5

1.0

Success

Lift

25

50

100

200

Demonstrations

PickPlaceCan

25

50

100

200

Demonstrations

NutAssemblySquare

bc

dac

iqlearn

csil

0

0.25

0.5

0.75

1.0

Steps (

1×10 6

)

0.50

0.75

1.00

Success

(

n

=200)

Lift

0

0.25

0.5

0.75

1.0

Steps (

1×10 6

)

PickPlaceCan

bc (csil)

csil

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].

6 Discussion

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.

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.

Acknowledgments and Disclosure of Funding

We wish to thank Jost Tobias Springenberg, Markus Wulfmeier, Todor Davchev, Ksenia Konyushkova, Abbas Abdolmaleki 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¨ ur K¨ unstliche Intelligenz (DFKI) an der Technischen Universit¨ at Darmstadt' of the Hessian Ministry of Science and Art.

References

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.

Stefan Schaal. Is imitation learning the route to humanoid robots? Trends in Cognitive Sciences , 1999.

D. A. Pomerleau. Efficient training of artificial neural networks for autonomous navigation. Neural Computation , 1991.

Rudolph E. Kalman. When is a linear control system optimal? Journal of Basic Engineering , 1964.

Andrew Y. Ng and Stuart J. Russell. Algorithms for inverse reinforcement learning. In International Conference on Machine Learning , 2000.

P. Abbeel and A. Y. Ng. Apprenticeship learning via inverse reinforcement learning. In International Conference on Machine Learning , 2004.

Martin L Puterman. Markov Decision Processes: Discrete Stochastic Dynamic Programming . John Wiley & Sons, 2014.

Richard S Sutton and Andrew G Barto. Reinforcement learning: An introduction . MIT press, 2018.

St´ ephane 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.

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.

Manu Orsini, Anton Raichuk, L´ eonard 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.

Rohit Jena, Changliu Liu, and Katia Sycara. Augmenting GAIL with BC for sample efficient imitation learning. In Conference on Robot Learning , 2020.

Jan Peters, Katharina M¨ ulling, and Yasemin Alt¨ un. Relative entropy policy search. In Proceedings of the AAAI Conference on Artificial Intelligence , 2010.

Roy Fox, Ari Pakman, and Naftali Tishby. Taming the noise in reinforcement learning via soft updates. In Conference on Uncertainty in Artificial Intelligence , 2016.

Herke Van Hoof, Gerhard Neumann, and Jan Peters. Non-parametric policy search with limited information loss. Journal of Machine Learning Research , 2017.

Abbas Abdolmaleki, Jost Tobias Springenberg, Yuval Tassa, Remi Munos, Nicolas Heess, and Martin Riedmiller. Maximum a Posteriori policy optimisation. In International Conference on Learning Representations , 2018.

Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: Offpolicy maximum entropy deep reinforcement learning with a stochastic actor. In International Conference on Machine Learning , 2018.

Matthew Fellows, Anuj Mahajan, Tim GJ Rudner, and Shimon Whiteson. VIREL: A variational inference framework for reinforcement learning. Advances in Neural Information Processing Systems , 2019.

Joan Bas-Serrano, Sebastian Curi, Andreas Krause, and Gergely Neu. Logistic Q -learning. In International Conference on Artificial Intelligence and Statistics , 2021.

Matteo Pirotta and Marcello Restelli. Inverse reinforcement learning through policy gradient minimization. In Proceedings of the AAAI Conference on Artificial Intelligence , 2016.

Giorgia Ramponi, Gianluca Drappo, and Marcello Restelli. Inverse reinforcement learning from a gradient-based learner. In Advances in Neural Information Processing Systems , 2020.

Yi-Feng Zhang, Fan-Ming Luo, and Yang Yu. Improve generated adversarial imitation learning with reward variance regularization. Machine Learning , 2022.

Siddharth Reddy, Anca D Dragan, and Sergey Levine. SQIL: Imitation learning via reinforcement learning with sparse rewards. In International Conference on Learning Representations , 2019.

Kiante Brantley, Wen Sun, and Mikael Henaff. Disagreement-regularized imitation learning. In International Conference on Learning Representations , 2020.

Nived Rajaraman, Lin Yang, Jiantao Jiao, and Kannan Ramchandran. Toward the fundamental limits of imitation learning. Advances in Neural Information Processing Systems , 2020.

Stuart Russell. Learning agents for uncertain environments. In Conference on Computational Learning Theory , 1998.

Umar Syed and Robert E. Schapire. A game-theoretic approach to apprenticeship learning. In Advances in Neural Information Processing Systems , 2007.

Umar Syed, M. Bowling, and R.E. Schapire. Apprenticeship learning using linear programming. In International Conference on Machine Learning , 2008.

Tian Xu, Ziniu Li, and Yang Yu. Error bounds of imitating policies and environments. Advances in Neural Information Processing Systems , 2020.

Brian D Ziebart. Modeling purposeful adaptive behavior with the principle of maximum causal entropy . Carnegie Mellon University, 2010.

Edwin T Jaynes. On the rationale of maximum-entropy methods. Proceedings of the IEEE , 1982.

Joseph Marino, Alexandre Pich´ e, Alessandro Davide Ialongo, and Yisong Yue. Iterative amortized policy optimization. In Advances in Neural Information Processing Systems , 2021.

David Barber. Bayesian reasoning and machine learning . Cambridge University Press, 2012.

Konrad Rawlik, Marc Toussaint, and Sethu Vijayakumar. On stochastic optimal control and reinforcement learning by approximate inference. In Robotics: Science and Systems , 2013.

Joe Watson and Jan Peters. Inferring smooth control: Monte Carlo posterior policy iteration with Gaussian processes. In Conference on Robot Learning , 2022.

Alberto Maria Metelli, Matteo Pirotta, and Marcello Restelli. Compatible reward inverse reinforcement learning. In Advances in Neural Information Processing Systems , 2017.

Jonathon Ho and Stefan Ermon. Generative adversarial imitation learning. In Advances in Neural Information Processing Systems , 2016.

Justin Fu, Katie Luo, and Sergey Levine. Learning robust rewards with adverserial inverse reinforcement learning. In International Conference on Learning Representations , 2018.

Seyed Kamyar Seyed Ghasemipour, Richard Zemel, and Shixiang Gu. A divergence minimization perspective on imitation learning methods. In Conference on Robot Learning , 2019.

Ilya Kostrikov, Kumar Krishna Agrawal, Debidatta Dwibedi, Sergey Levine, and Jonathan Tompson. Discriminator-actor-critic: Addressing sample inefficiency and reward bias in adversarial imitation learning. In International Conference on Learning Representations , 2019.

Ilya Kostrikov, Ofir Nachum, and Jonathan Tompson. Imitation learning via off-policy distribution matching. In International Conference on Learning Representations , 2020.

Divyansh Garg, Shuvam Chakraborty, Chris Cundy, Jiaming Song, and Stefano Ermon. IQlearn: Inverse softQ learning for imitation. In Advances in Neural Information Processing Systems , 2021.

Luca Viano, Angeliki Kamoutsi, Gergely Neu, Igor Krawczuk, and Volkan Cevher. Proximal point imitation learning. In Advances in Neural Information Processing Systems , 2022.

Robert Dadashi, Leonard Hussenot, Matthieu Geist, and Olivier Pietquin. Primal Wasserstein imitation learning. In International Conference on Learning Representations , 2021.

Abdeslam Boularias, Jens Kober, and Jan Peters. Relative entropy inverse reinforcement learning. In International Conference on Artificial Intelligence and Statistics , 2011.

Markus Wulfmeier, Peter Ondruska, and Ingmar Posner. Maximum entropy deep inverse reinforcement learning. arXiv preprint arXiv:1507.04888 , 2015.

Oleg Arenz, Hany Abdulsamad, and Gerhard Neumann. Optimal control and inverse optimal control by distribution matching. In IEEE/RSJ International Conference on Intelligent Robots and Systems , 2016.

Chelsea Finn, Sergey Levine, and Pieter Abbeel. Guided cost learning: Deep inverse optimal control via policy optimization. In International Conference on Machine Learning , 2016.

Tianwei Ni, Harshit Sikchi, Yufei Wang, Tejus Gupta, Lisa Lee, and Ben Eysenbach. f -IRL: Inverse reinforcement learning via state marginal matching. In Conference on Robot Learning , 2020.

Gergeley Neu and Csaba Szepesv´ ari. Apprenticeship learning using inverse reinforcement learning and gradient methods. In Conference on Uncertainty in Artificial Intelligence , 2007.

Paul Barde, Julien Roy, Wonseok Jeon, Joelle Pineau, Chris Pal, and Derek Nowrouzezahrai. Adversarial soft advantage fitting: Imitation learning without policy optimization. In Advances in Neural Information Processing Systems , 2020.

Haoran Xu, Xianyuan Zhan, Honglei Yin, and Huiling Qin. Discriminator-weighted offline imitation learning from suboptimal demonstrations. In International Conference on Machine Learning , 2022.

Alexis Jacq, Matthieu Geist, Ana Paiva, and Olivier Pietquin. Learning from a learner. In International Conference on Machine Learning , 2019.

Haoyang Cao, Samuel Cohen, and Lukasz Szpruch. Identifiability in inverse reinforcement learning. In Advances in Neural Information Processing Systems , 2021.

Andrew Y. Ng, Daishi Harada, and Stuart J. Russell. Policy invariance under reward transformations: Theory and application to reward shaping. In International Conference on Machine Learning , 1999.

Jan van Campenhout and Thomas Cover. Maximum entropy and conditional probability. IEEE Transactions on Information Theory , 1981.

Benjamin Guedj. A primer on PAC-Bayesian learning. In Proceedings of the second congress of the French Mathematical Society , 2019.

Pierre Alquier, James Ridgway, and Nicolas Chopin. On the properties of variational approximations of Gibbs posteriors. Journal of Machine Learning Research , 2016.

Jeremias Knoblauch, Jack Jewson, and Theodoros Damoulas. An optimization-centric view on Bayes' rule: Reviewing and generalizing variational inference. Journal of Machine Learning Research , 2022.

Hanbit Oh, Hikaru Sasaki, Brendan Michael, and Takamitsu Matsubara. Bayesian disturbance injection: Robust imitation learning of flexible policies for robot manipulation. Neural Networks , 2023.

Tim GJ Rudner, Cong Lu, Michael A Osborne, Yarin Gal, and Yee Teh. On pathologies in KL-regularized reinforcement learning from expert demonstrations. Advances in Neural Information Processing Systems , 2021.

Carl E. Rasmussen and Chris Williams. Gaussian Processes for Machine Learning . MIT Press, 2006.

Quoc V. Le, Alex J. Smola, and St´ ephane Canu. Heteroscedastic Gaussian process regression. In International Conference on Machine Learning , 2005.

Martin Jankowiak, Geoff Pleiss, and Jacob Gardner. Parametric Gaussian process regressors. In International Conference on Machine Learning , 2020.

Justin Fu, Katie Luo, and Sergey Levine. Learning robust rewards with adverserial inverse reinforcement learning. In International Conference on Learning Representations , 2018.

Matthieu Geist, Bruno Scherrer, and Olivier Pietquin. A theory of regularized Markov decision processes. In International Conference on Machine Learning , 2019.

David Roxbee Cox and Hilton David Miller. The theory of stochastic processes . CRC press, 1977.

Lassi Meronen, Martin Trapp, and Arno Solin. Periodic activation functions induce stationarity. In Advances in Neural Information Processing Systems , 2021.

Noah Siegel, Jost Tobias Springenberg, Felix Berkenkamp, Abbas Abdolmaleki, Michael Neunert, Thomas Lampe, Roland Hafner, Nicolas Heess, and Martin Riedmiller. Keep doing what worked: Behavior modelling priors for offline reinforcement learning. In International Conference on Learning Representations , 2020.

Maximilian Seitzer, Arash Tavakoli, Dimitrije Antic, and Georg Martius. On the pitfalls of heteroscedastic uncertainty estimation with probabilistic neural networks. In International Conference on Learning Representations , 2022.

Andrew Stirn, Hans-Hermann Wessels, Megan Schertzer, Laura Pereira, Neville E. Sanjana, and David A. Knowles. Faithful heteroscedastic regression with neural networks. In International Conference on Artificial Intelligence and Statistics , 2023.

Emanuel Todorov, Tom Erez, and Yuval Tassa. MuJoCo: A physics engine for model-based control. In IEEE International Conference on Intelligent Robots and Systems , 2012.

Greg Brockman, Vicki Cheung, Ludwig Pettersson, Jonas Schneider, John Schulman, Jie Tang, and Wojciech Zaremba. OpenAI Gym. arXiv preprint arXiv:1606.01540 , 2016.

Aravind Rajeswaran, Vikash Kumar, Abhishek Gupta, Giulia Vezzani, John Schulman, Emanuel Todorov, and Sergey Levine. Learning complex dexterous manipulation with deep reinforcement learning and demonstrations. In Robotics: Science and Systems , 2018.

Justin Fu, Aviral Kumar, Ofir Nachum, George Tucker, and Sergey Levine. D4rl: Datasets for deep data-driven reinforcement learning. arXiv preprint arXiv:2004.07219 , 2020.

Matthew W Hoffman, Bobak Shahriari, John Aslanides, Gabriel Barth-Maron, Nikola Momchev, Danila Sinopalnikov, Piotr Sta´ nczyk, Sabela Ramos, Anton Raichuk, Damien Vincent, et al. Acme: A research framework for distributed reinforcement learning. arXiv preprint arXiv:2006.00979 , 2020.

Sabela Ramos, Sertan Girgin, L´ eonard Hussenot, Damien Vincent, Hanna Yakubovich, Daniel Toyama, Anita Gergely, Piotr Stanczyk, Raphael Marinier, Jeremiah Harmsen, et al. Rlds: An ecosystem to generate, share and use datasets in reinforcement learning. arXiv preprint arXiv:2111.02767 , 2021.

Mel Vecerik, Todd Hester, Jonathan Scholz, Fumin Wang, Olivier Pietquin, Bilal Piot, Nicolas Heess, Thomas Roth¨ orl, Thomas Lampe, and Martin Riedmiller. Leveraging demonstrations for deep reinforcement learning on robotics problems with sparse rewards. arXiv preprint arXiv:1707.08817 , 2017.

Oleg Arenz and Gerhard Neumann. Non-adversarial imitation learning and its connections to adversarial methods. arxiv preprint arxiv:2008.03525 , 2020.

Li Ziniu, Xu Tian, Yu Yang, and Luo Zhi-Quan. Rethinking ValueDice - Does it really improve performance? ICLR Blog Track, https://iclr-blog-track.github.io/2022/03/25/ rethinking-valuedice/ , 2022.

Aviral Kumar, Aurick Zhou, George Tucker, and Sergey Levine. Conservative Q-learning for offline reinforcement learning. Advances in Neural Information Processing Systems , 2020.

Yecheng Ma, Andrew Shen, Dinesh Jayaraman, and Osbert Bastani. Versatile offline imitation from observations and examples via regularized state-occupancy matching. In International Conference on Machine Learning , 2022.

Geon-Hyeong Kim, Seokin Seo, Jongmin Lee, Wonseok Jeon, HyeongJoo Hwang, Hongseok Yang, and Kee-Eung Kim. DemoDICE: Offline imitation learning with supplementary imperfect demonstrations. In International Conference on Learning Representations , 2022.

Ajay Mandlekar, Danfei Xu, Josiah Wong, Soroush Nasiriany, Chen Wang, Rohun Kulkarni, Li Fei-Fei, Silvio Savarese, Yuke Zhu, and Roberto Mart´ ın-Mart´ ın. What matters in learning from offline human demonstrations for robot manipulation. In Conference on Robot Learning , 2021.

Konrad Zołna, Alexander Novikov, Ksenia Konyushkova, Caglar Gulcehre, Ziyu Wang, Yusuf Aytar, Misha Denil, Nando de Freitas, and Scott Reed. Offline learning from demonstrations and unlabeled experience. In Offline Reinforcement Learning Workshop, NeurIPS , 2020.

Marc Peter Deisenroth, Gerhard Neumann, and Jan Peters. A survey on policy search for robotics. Foundations and Trends in Robotics , 2013.

Firas Al-Hafez, Davide Tateo, Oleg Arenz, Guoping Zhao, and Jan Peters. LS-IQ: Implicit reward regularization for inverse reinforcement learning. In International Conference on Learning Representations , 2023.

Andrew Szot, Amy Zhang, Dhruv Batra, Zsolt Kira, and Franziska Meier. BC-IRL: Learning generalizable reward functions from demonstrations. In International Conference on Learning Representations , 2023.

Aleksandar Taranovic, Andras Gabor Kupcsik, Niklas Freymuth, and Gerhard Neumann. Adversarial imitation learning with preferences. In International Conference on Learning Representations , 2022.

Gokul Swamy, Nived Rajaraman, Matt Peng, Sanjiban Choudhury, J Bagnell, Steven Z Wu, Jiantao Jiao, and Kannan Ramchandran. Minimax optimal online imitation learning via replay estimation. Advances in Neural Information Processing Systems , 2022.

Fumihiro Sasaki and Ryota Yamashina. Behavioral cloning from noisy demonstrations. In International Conference on Learning Representations , 2021.

Deepak Ramachandran and Eyal Amir. Bayesian inverse reinforcement learning. In International Joint Conference on Artifical Intelligence , 2007.

Jaedeug Choi and Kee-Eung Kim. MAP inference for Bayesian inverse reinforcement learning. Advances in Neural Information Processing Systems , 2011.

Giuseppe Da Prato and Jerzy Zabczyk. Stochastic equations in infinite dimensions . Cambridge University Press, 2014.

Chris Chatfield. The analysis of time series: An introduction . Chapman and Hall, 1989.

Jeremiah Liu, Zi Lin, Shreyas Padhy, Dustin Tran, Tania Bedrax Weiss, and Balaji Lakshminarayanan. Simple and principled uncertainty estimation with deterministic deep learning via distance awareness. In Advances in Neural Information Processing Systems , 2020.

John Schulman. Approximating KL divergence. http://joschu.net/blog/kl-approx. html , 2020.

Shengyang Sun, Guodong Zhang, Jiaxin Shi, and Roger Grosse. Functional variational Bayesian neural networks. In International Conference on Learning Representations , 2019.

Joe Watson, Jihao Andreas Lin, Pascal Klink, and Jan Peters. Neural linear models with functional Gaussian process priors. In Third Symposium on Advances in Approximate Bayesian Inference , 2021.

Div99 . IQ-Learn implementation. https://github.com/Div99/IQ-Learn , 2021.

Yihong Wu. Lecture notes for ECE598YW: Information-theoretic methods in high-dimensional statistics. Yale University, 2016.

James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman-Milne, and Qiao Zhang. JAX: Composable transformations of Python+NumPy programs, 2018. URL http://github.com/google/jax .

Igor Babuschkin, Kate Baumli, Alison Bell, Surya Bhupatiraju, Jake Bruce, Peter Buchlovsky, David Budden, Trevor Cai, Aidan Clark, Ivo Danihelka, Antoine Dedieu, Claudio Fantacci, Jonathan Godwin, Chris Jones, Ross Hemsley, Tom Hennigan, Matteo Hessel, Shaobo Hou, Steven Kapturowski, Thomas Keck, Iurii Kemaev, Michael King, Markus Kunesch, Lena Martens, Hamza Merzic, Vladimir Mikulik, Tamara Norman, George Papamakarios, John Quan, Roman Ring, Francisco Ruiz, Alvaro Sanchez, Laurent Sartran, Rosalia Schneider, Eren Sezener, Stephen Spencer, Srivatsan Srinivasan, Miloˇ s Stanojevi´ c, Wojciech Stokowiec, Luyu Wang, Guangyao Zhou, and Fabio Viola. The DeepMind JAX Ecosystem, 2020. URL http://github.com/deepmind .

lviano . Proximal point imitation learning implementation. https://github.com/lviano/ p2il , 2022.

geon-hyeong . DemoDICE implementation. https://github.com/geon-hyeong/ imitation-dice , 2022.

JasonMa2016 . SMODICE implementation. https://github.com/JasonMa2016/ SMODICE , 2022.

Lasse Espeholt, Hubert Soyer, Remi Munos, Karen Simonyan, Vlad Mnih, Tom Ward, Yotam Doron, Vlad Firoiu, Tim Harley, Iain Dunning, et al. Impala: Scalable distributed deep-RL with importance weighted actor-learner architectures. In International conference on machine learning , 2018.

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 onpolicy 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] highlight 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, simplififying 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 maximimum 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 ± 1 rather than assume that the likelihood values are meaningful.

Taranovic et al. [89] derive an AIRL formalation 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,

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 = w ⊤ ϕ ( x ) , w ∼ N ( µ , Σ ) , through their kernel function K ( x , x ′ ) = ϕ ( x ) ⊤ ϕ ( x ′ ) , the feature inner product. For stationary kernels, the kernel depends only on a relative shift r ∈ R d between inputs, i.e. K ( r ) = ϕ ( x ) ⊤ ϕ ( x + r ) . By combining the kernel with the weight prior of the linear model, we can define the covariance function, e.g. C ( r ) = ϕ ( x ) ⊤ Σ ϕ ( 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 ( ω ) , and it is a Fourier dual of the covariance function if the necessary conditions apply, known as the Wiener-Khinchin theorem [95],

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ω ⊤ 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

s

a

Heteroscedastic

MLP

(NLLH)

Train MSE = 0.11

s

Heteroscedastic

MLP

(Faithful)

Train MSE = 0.05

s

Stationary

Heteroscedastic

MLP (NLLH)

Train MSE = 0.04

s

Stationary

Heteroscedastic

MLP (Faithful)

Train MSE = 0.02

is dictated by 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 Studentt or Cauchy could be used if 'rougher' processes are desired. In summary, the last layer activations take the form f p ( Wϕ ( x )) , where f p is the periodic activation function of choice, W are weights sampled from the distribution of choice and ϕ ( 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.

Acrucial 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

where sg ( · ) 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

If the behavior policy π is q θ ( a | s ) , we can view the second term as a Monte Carlo estimation of the relative KL divergence between q and p . This estimator can suffer greatly from bias, as the KL should always be non-negative. To counteract this, we construct an unbiased, positively-constrained estimator following Schulman [97], using E x ∼ q ( · ) [ R ( x ) -1 ] = ∫ p ( x ) d x = 1 and log( x ) ≤ x -1 ,

This estimator results in replacing the second term of Equation 11 with

This objective is maximized concurrently with policy evaluation, using samples from the replay buffer. This reward refinement is similar to behavioral cloning via function-space variational inference (e.g., [98, 99]), enforcing the prior using the prior predictive distribution rather than the prior weight distribution.

In practice, the samples are taken from a replay buffer of online and offline data, so the action samples are taken from a different distribution to q θ . However, we still found the adjusted objective effective for reward regularization. Without it, the saddle-point refinement was less stable. Appendix G illustrates this refinement in a simple setting. Figure 45 ablates this component, where its effect is most beneficial when there are fewer demonstrations.

F Bounding the coherent reward

The coherent reward α (log q ( a | s ) -log p ( a | s )) can be upper bounded if p ( a | s ) > 0 whenever q ( a | s ) > 0 . The bound is policy dependent. We use the bound in the continuous setting, where we use a tanh-transformed Gaussian policy. We add a small bias term σ 2 min to the predictive variance in order to define the Gaussian likelihood upper bound. Inconveniently, the tanh change of variables term in the log-likelihood -∑ d a i =1 log(1 -tanh 2 ( u )) [17] has no upper bound in theory, as the 'latent' action u ∈ [ -∞ , ∞ ] , but in the Acme implementation this log-likelihood is bounded for numerical stability due to the inverse tanh operation. For action dimension d a and uniform prior across action range [ -1 , 1] d a and α = d -1 a , the upper bound is r ( s , a ) ≤ 1 d a ( -0 . 5 d a log 2 πσ 2 min + c + d a log 2) = -0 . 5 log πσ 2 min / 2 + ˜ c , where c, ˜ c depends on the tanh clipping term.

G Visualizing reward refinement

To provide intuition on Theorem 2 and connect it to reward refinement (Appendix E), we build on the regression example from Figure 3 by viewing it as a continuous contextual bandit problem, i.e., a single-state MDP with no dynamics. Outside of the demonstration state distribution, we desire that the reward is uniformly zero, as we have no data to inform action preference. Figure 10 shows the result of applying CSIL and reward refinement to this MDP, where the state distribution is uniform. The 'desired' BC fit exhibits a coherent reward with the desired qualities, and the HETSTAT policy approximates this desired policy. The heteroscedastic MLP exhibits undesirable out-of-distribution (OOD) behavior, where arbitrary regions have strongly positive or negative rewards. Reward refinement improves the coherent reward of the heteroscedastic MLP to be more uniform OOD, which in turn makes the regression fit look more like a stationary policy. This result reinforces Theorem 2, which shows that KL-regularized BC is connected to a lower bound on an entropy-regularized game-theoretic IRL objective that uses the coherent reward.

Desired

Heteroscedastic MLP

Stationary Heteroscedastic MLP

a

Initial regression

Initial regression

a

Initial coherent reward

Initial coherent reward

a

Refined regression

Refined regression

s

a

s

Refined coherent reward

s

Refined coherent reward

-

9

.

0

-

7

.

5

-

6

.

0

-

4

.

5

-

3

.

0

-

1

.

5

0

.

0

1

.

5

3

.

0

H Critic regularization

For the function approximation setting, CSIL requires an auxiliary loss term during learning to incorporate coherency. The full critic optimization problem is min ϕ J Q ( ϕ ) for objective

where Q ∗ denotes the target Q values and B denotes the replay buffer, that combines demonstrations D and additional samples. The motivation is to shape the critic such that the demonstration actions are first-order optimal by minimizing the squared Frobenius norm of the action Jacobian.

IQLearn and PPIL adopt similar, but less explicit, regularization. Equation 10 of Garg et al. [42] is

where V π ϕ ( s ) = Q ( s , a ′ ) , a ′ ∼ π ( · | s ) and f is the concave function defining the implicit reward regularization. In the implementation [100], the ' v0 ' objective variant replaces the initial distribution µ 0 with the demonstration state distribution. This trick is deployed to improve regularization without affecting optimality (Section 5.3, Kostrikov et al. [41]). If the policy matches the expert actions with its mode, this term shapes a local maximum by minimizing the critic at action samples around the

Q

Jacobian regularization

sample regularization

∂Q/∂a

a

a

Algorithm 2: Off-policy coherent soft imitation learning with function approximation

expert. However, if the policy does not accurately match the expert actions, the effect of this term is harder to reason about. In Figure 11, we illustrate these two forms from critic regularization on a toy example, demonstrating that both can produce local maximums in the critic.

I Practical algorithm with function approximation

For deep imitation learning, CSIL adopts SAC for the SPI step in Algorithm 1, summarized in Algorithm 2. The critic objective for both pre-training and policy evaluation is the squared Bellman error. The policy objective max θ J π ( θ ) samples both the demonstrations and replay buffer,

The policy and demonstration data are combined equally. We did not experiment with tuning this ratio. Compared to SAC, we found we did not have to train multiple critics, but we did need target networks for the critics to stabilize learning. The critic and reward objectives were also evaluated on the policy and expert data each update. Moreover, like IQLearn and PPIL, we kept β constant and did not use an adaptive strategy.

J A divergence minimization perspective

Many imitation learning algorithms, such as GAIL and IQLearn, are derived from minimizing a divergence between the agent's state-action distribution and the expert's. While CSIL is derived through policy inversion, it also has a divergence minimization perspective, albeit with two minimization steps rather than one.

The first step involves recognising that maximum likelihood BC fit minimizes the forward KL divergence between the demonstration distribution p D ( s , a ) = p D ( a | s ) p D ( s ) and the policy,

Using the forward KL means that the unknown log-likelihood of the data distribution is not required, and the policy can be trained using supervised learning and without environment interaction.

The second step involves rearranging the KL-regularized RL objective into two KL divergences,

Note that since β ≤ α , then (1 -β/α ) ≥ 0 . The left-hand term minimizes the reverse KL between the agent and the expert, using the BC policy as a proxy. This is beneficial because the reverse conditional KL requires evaluating policy q θ on the environment. Moreover, the reverse KL requires log-likelihoods of the target distribution, which is why the BC policy is used as a proxy of the demonstration distribution. The right-hand term incorporates the coherent reward and encourages the policy to stay in states where the policy differs from the prior. As q θ should be close to the BC policy due to the left-hand term, the coherent reward shaping from Definition 1 should hold.

In this formulation of the objective, it is also clear to see that if β = α , no RL finetuning will occur, and q θ will simply mimic the BC policy q θ 1 due to the left-hand term.

K Theoretical results

This section provides derivations and proofs for the theoretical results in the main text.

K.1 Relative entropy regularization for reinforcement learning

This section re-derives KL-regularized RL and ME-IRL in a shared notation for the reader's reference.

Reinforcement learning. We consider the discounted, infinite-horizon setting with a hard KL constraint on the policy update with bound ϵ ,

This constrained objective is transcribed into the following Lagrangian objective following Van Hoof et al. [15], using the discounted stationary joint distribution d ( s , a ) = q ( a | s ) ν ( s ) = (1 -γ ) ρ ( s , a ) , where λ 2 is a function for an infinite-dimensional constraint. Solving for ∂ L ( q, λ 1 , λ 2 , λ 3 ) /∂d = 0 ,

with α = (1 -γ ) λ 3 and λ 2 ( s ) = (1 -γ ) V ( s ) , where V can be interpreted as the soft value function,

To obtain the more convenient lower bound of the soft Bellman equation, we combine importance sampling and Jensen's inequality to the soft Bellman equation

This derivation is for the 'hard' KL constraint. In many implementations, including ours, the constraint is 'softened' where λ 3 and, therefore, α is a fixed hyperparameter.

Inverse reinforcement learning. As mentioned in Section 2 and Definition 2, ME-IRL is closely related to KL-regularize RL, but with the objective and a constraint reversed,

The constrained optimization is, therefore, transcribed into a similar Lagrangian,

The consequence of this change is that the likelihood temperature is now implicit in the reward model r ( s , a ) = (1 -γ ) -1 λ ⊤ 3 ϕ ( s , a ) due to the Lagrange multiplier weights, this reward model is used in the soft Bellman equation instead of the true reward. The reward model is jointly optimized to minimize the apprenticeship error by matching the feature expectation.

K.2 Proof for Theorem 1

Theorem 1. (KL-regularized policy improvement inversion). Let p and q α be the prior and pseudoposterior policy given by posterior policy iteration (Equation 4). The critic can be expressed as

Substituting into the KL-regularized Bellman equation lower-bound from Equation 2,

The V α ( s ) term is the 'soft' value function. We assume q α ( a | s ) = 0 whenever p ( a | s ) = 0 .

Proof. Equation 6 is derived by rearranging the posterior policy update in Equation 4. Equation 7 in derived by substituting the critic expression from Equation 6 into the KL-regularized Bellman equation (Equation 3),

K.3 Proof for Theorem 2

Theorem 2 relies on the data processing inequality (Lemma 3) to connect the function-space KL seen in CSIL with the weight-space KL used in regularized BC.

.

Proof. This result arises from the origin of the coherent reward from policy inversion and reward shaping (Theorem 1, Lemma 1). Moreover, considering the entropy-augmented reward r α, π ( s , a ) = r ( s , a ) -α (log π ( a | s ) -log p ( a | s ) , when r is the coherent reward ˜ r and π = q , ˜ r α, q ( s , a ) = 0 ∀ s ∈ S , a ∈ A so q ( a | s ) is a solution to a soft policy iteration.

Assumption 1. (Realizability) The optimal policy can be expressed by the parametric policy q θ ( a | s ) Lemma 2. (Coherent policy optimality). The optimal policy for a KL -regularized RL problem with coherent reward ˜ r ( s , a ) = α (log q ( a | s ) -log p ( a | s )) and temperature α is q ( a | s ) .

Lemma3. (Data processing inequality, Wu [101], Theorem 4.1). For two stochastic processes parameterized with finite random variables w ∼ p ( · ) , w ∈W and shared hypothesis spaces p ( y | x , w ) , the conditional f -divergence in function space at points X ∈ X L , L ∈ Z + . is upper bounded by the f -divergence in parameter space,

Theorem 2. (Coherent inverse reinforcement learning as KL -regularized behavioral cloning). A KL -regularized game-theoretic IRL objective, with policy q θ ( a | s ) and coherent reward parameterization r θ ( s , a ) = α (log q θ ( a | s ) -log p ( a | s )) where α ≥ 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 ( a | s , w ) and finite parameters w , where λ = ( α -β ) /α and E s , a ∼D [log p ( a | s )] is constant. The regression objective bounds the IRL objective for a worst-case on-policy state distribution µ q ( s ) , which motivates its scaling through β . If D has sufficient coverage, no RL finetuning or KL regularization is required, so β = α and λ = 0 . If D does not have sufficient coverage, then let β < α so λ > 0 to regularize the BC fit and finetune the policy with RL accordingly with additional soft policy iteration steps.

Proof. We begin with the KL-regularized game-theoretic IRL objective (c.f. Equation 1 [17]),

Substituting the coherent reward from Theorem 1, the objective becomes

Due to the realizablity assumption (Assumption 1), the optimal policy π θ ∗ can be obtained through maximum likelihood estimation θ ∗ = arg max θ E s , a ∼D [log q θ ( s , a )] . With this assumption, the lefthand 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 π with the maximum likelihood solution of q θ , resulting in a simpler objective ˜ J in terms of only θ , which matches the objective J exactly when β = α and π = q θ (Lemma 2). We consider ˜ J with a small negative perturbation β = α -δα , δα ≥ 0 to examine the effect of the right-term term,

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 λ = δα/α . The resulting lower-bound objective is proportional to the regularized behavioral cloning objective proposed in Equation 8,

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 acme , again using jax and its ecosystem [103]. One feature of acme 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 acme . 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 d4rl rather than Orsini et al. [11], but the returns were normalized to the d4rl returns for consistency. We note that there is a non-negligible discrepancy between the d4rl 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 acme and was found to improve performance. For the inverse temperature α , 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

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.

Value function

Nominal stationary distribution

Windy stationary distribution

Value function

Nominal stationary distribution

Windy stationary distribution

Value function

Nominal stationary distribution

Windy stationary distribution

Value function

Nominal stationary distribution

Windy stationary distribution

Nominal stationary distribution

Windy stationary distribution

Value function

Nominal stationary distribution

Windy stationary distribution

Value function

Nominal stationary distribution

Windy stationary distribution

Value function

Nominal stationary distribution

Windy stationary distribution

Value function

Nominal stationary distribution

Windy stationary distribution

Value function

Nominal stationary distribution

Windy stationary distribution

Value function

Nominal stationary distribution

Windy stationary distribution

Value function

Nominal stationary distribution

Windy stationary distribution

Nominal stationary distribution

Windy stationary distribution

Value function

Nominal stationary distribution

Windy stationary distribution

Value function

Nominal stationary distribution

Windy stationary distribution

Value function

Nominal stationary distribution

Windy stationary distribution

0

1

Return

(

n

=1)

HalfCheetah-v2

Hopper-v2

0

1

Return

(

n

=3)

0

1

Return

(

n

=10)

0

0.25

0.5

0.75

1

Steps (

1×10 6

)

0

1

Return

(

n

=30)

bc

sqil

dac

iqlearn

ppil

pwil

csil

0

0.25

0.5

0.75

1

Steps (

1×10 6

)

0

0.25

0.5

0.75

1

Steps (

1×10 6

)

0

0.25

0.5

0.75

1

Steps (

1×10 6

)

0

1

Return

(

n

=1)

door-v0

hammer-v0

0

1

Return

(

n

=3)

0

1

Return

(

n

=10)

0

0.25

0.5

0.75

1

Steps (

1×10 6

)

0

1

Return

(

n

=30)

bc

sacfd

sqil

dac

iqlearn

ppil

pwil

csil

0

0.25

0.5

0.75

1

Steps (

1×10 6

)

0

1

Return

(

n

=1)

HalfCheetah-v2

Hopper-v2

0

1

Return

(

n

=3)

0

1

Return

(

n

=10)

0

0.25

0.5

0.75

1

Steps (

1×10 6

)

0

1

Return

(

n

=30)

cql

bc

iqlearn

ppil

csil

0

0.25

0.5

0.75

1

Steps (

1×10 6

)

0

0.25

0.5

0.75

1

Steps (

1×10 6

)

0

0.25

0.5

0.75

1

Steps (

1×10 6

)

0.0

0.5

1.0

Return

(

n

=1)

Humanoid-v2

0.0

0.5

1.0

Return

(

n

=3)

0.0

0.5

1.0

Return

(

n

=10)

0

0.25

0.5

0.75

1

Steps (

1×10 6

)

0.0

0.5

1.0

Return

(

n

=30)

bc

pwil

sqil

dac

iqlearn

ppil

csil

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.

0.0

0.5

1.0

Return

(

n

=1)

HalfCheetah-v2

Hopper-v2

0.0

0.5

1.0

Return

(

n

=3)

0.0

0.5

1.0

Return

(

n

=10)

0

0.25

0.5

0.75

1

Steps (

1×10 6

)

0.0

0.5

1.0

Return

(

n

=30)

0

0.25

0.5

0.75

1

Steps (

1×10 6

)

0

0.25

0.5

0.75

1

Steps (

1×10 6

)

0

0.25

0.5

0.75

1

Steps (

1×10 6

)

bc (csil)

csil

0.0

0.5

1.0

Return

(

n

=1)

door-v0

hammer-v0

0.0

0.5

1.0

Return

(

n

=3)

0.0

0.5

1.0

Return

(

n

=10)

0

0.25

0.5

0.75

1

Steps (

1×10 6

)

0.0

0.5

1.0

Return

(

n

=30)

0

0.25

0.5

0.75

1

Steps (

1×10 6

)

bc (csil)

csil

0.0

0.5

1.0

Return

(

n

=1)

HalfCheetah-v2

Hopper-v2

0.0

0.5

1.0

Return

(

n

=3)

0.0

0.5

1.0

Return

(

n

=10)

0

0.25

0.5

0.75

1

Steps (

1×10 6

)

0.0

0.5

1.0

Return

(

n

=30)

0

0.25

0.5

0.75

1

Steps (

1×10 6

)

0

0.25

0.5

0.75

1

Steps (

1×10 6

)

0

0.25

0.5

0.75

1

Steps (

1×10 6

)

bc (csil)

csil

0

1

2

3

expert_reward

0.5

1.0

1.5

2.0

door-v0

0

100

200

300

400

500

expert_q

0

100

200

300

400

500

0

0.25

0.5

0.75

1

Steps (

1×10 6

)

0.0

0.2

0.4

expert_mse

0

0.25

0.5

0.75

1

Steps (

1×10 6

)

0.0

0.2

0.4

iqlearn ppil

csil

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.

0.0

0.5

1.0

Success

(

n

=25)

Lift

PickPlaceCan

NutAssemblySquare

0.0

0.5

1.0

Success

(

n

=50)

0.0

0.5

1.0

Success

(

n

=100)

0

0.25

0.5

0.75

1.0

Steps (

1×10 6

)

0.0

0.5

1.0

Success

(

n

=200)

0

0.25

0.5

0.75

1.0

Steps (

1×10 6

)

bc

sacfd

dac

iqlearn

csil

0

0.25

0.5

0.75

1.0

Steps (

1×10 6

)

0

0.25

0.5

0.75

1.0

Steps (

1×10 6

)

0.0

0.5

1.0

Success

(

n

=200)

Lift

0

0.25

0.5

0.75

1.0

Steps (

1×10 6

)

PickPlaceCan

bc (csil)

csil

0

0.25

0.5

0.75

1.0

Steps (

1×10 6

)

NutAssemblySquare

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.

0

1

Return

(

n

=1)

HalfCheetah-v2

Hopper-v2

0

1

Return

(

n

=3)

0

1

Return

(

n

=10)

0

0.25

Steps (

1×10 6

)

0

1

Return

(

n

=30)

iqlearn (w/ pretraining)

ppil (w/ pretraining)

iqlearn

ppil

0

0.25

Steps (

1×10 6

)

0

0.25

Steps (

1×10 6

)

0

0.25

Steps (

1×10 6

)

0.0

0.5

1.0

Return

(

n

=1)

door-v0

hammer-v0

0.0

0.5

1.0

Return

(

n

=3)

0.0

0.5

1.0

Return

(

n

=10)

0

0.25

0.5

0.75

1.0

Steps (

1×10 6

)

0.0

0.5

1.0

Return

(

n

=30)

iqlearn (w/ bc prior)

ppil (w/ bc prior)

iqlearn

ppil

0

0.25

0.5

0.75

1.0

Steps (

1×10 6

)

0

1

Return

(

n

=1)

HalfCheetah-v2

Hopper-v2

0

1

Return

(

n

=3)

0

1

Return

(

n

=10)

0

0.25

0.5

0.75

1.0

Steps (

1×10 6

)

0

1

Return

(

n

=30)

iqlearn (w/ bc prior)

ppil (w/ bc prior)

iqlearn

ppil

0

0.25

0.5

0.75

1.0

Steps (

1×10 6

)

0

0.25

0.5

0.75

1.0

Steps (

1×10 6

)

0

0.25

0.5

0.75

1.0

Steps (

1×10 6

)

0

1

Return

(

n

=1)

HalfCheetah-v2

Hopper-v2

0

1

Return

(

n

=3)

0

1

Return

(

n

=10)

0

0.25

0.5

0.75

1.0

Steps (

1×10 6

)

0

1

Return

(

n

=30)

iqlearn (w/ hetstat policy)

ppil (w/ hetstat policy)

iqlearn

ppil

0

0.25

0.5

0.75

1.0

Steps (

1×10 6

)

0

0.25

0.5

0.75

1.0

Steps (

1×10 6

)

0

0.25

0.5

0.75

1.0

Steps (

1×10 6

)

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.

0.0

0.5

1.0

Return

(

n

=1)

door-v0

hammer-v0

0.0

0.5

1.0

Return

(

n

=3)

0.0

0.5

1.0

Return

(

n

=10)

0

0.25

0.5

0.75

1.0

Steps (

1×10 6

)

0.0

0.5

1.0

Return

(

n

=30)

iqlearn (w/ pretrained hetstat prior)

ppil (w/ pretrained hetstat prior)

iqlearn

ppil

0

0.25

0.5

0.75

1.0

Steps (

1×10 6

)

0

1

Return

(

n

=1)

HalfCheetah-v2

Hopper-v2

0

1

Return

(

n

=3)

0

1

Return

(

n

=10)

0

0.25

0.5

0.75

1.0

Steps (

1×10 6

)

0

1

Return

(

n

=30)

iqlearn (w/ pretrained hetstat prior)

ppil (w/ pretrained hetstat prior)

iqlearn

ppil

0

0.25

0.5

0.75

1.0

Steps (

1×10 6

)

0

0.25

0.5

0.75

1.0

Steps (

1×10 6

)

0

0.25

0.5

0.75

1.0

Steps (

1×10 6

)

0.0

0.5

1.0

Return

(

n

=1)

HalfCheetah-v2

Hopper-v2

0.0

0.5

1.0

Return

(

n

=3)

0.0

0.5

1.0

Return

(

n

=10)

0

0.25

Steps (

1×10 6

)

0.0

0.5

1.0

Return

(

n

=30)

csil (w/o pretrained critic)

csil

0

0.25

Steps (

1×10 6

)

0

0.25

Steps (

1×10 6

)

0

0.25

Steps (

1×10 6

)

0.0

0.5

1.0

Return

(

n

=1)

HalfCheetah-v2

Hopper-v2

0.0

0.5

1.0

Return

(

n

=3)

0.0

0.5

1.0

Return

(

n

=10)

0

0.25

0.5

0.75

1.0

Steps (

1×10 6

)

0.0

0.5

1.0

Return

(

n

=30)

csil (w/o reward finetuning)

csil

0

0.25

0.5

0.75

1.0

Steps (

1×10 6

)

0

0.25

0.5

0.75

1.0

Steps (

1×10 6

)

0

0.25

0.5

0.75

1.0

Steps (

1×10 6

)

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.

N.8 Stationary and non-stationary policies

One important ablation for CSIL is evaluating the performance with standard MLP policies like those used by the baselines. We found that this setting was not numerically stable enough to complete an ablation experiment. Due to the spurious behavior (e.g., as shown in Figure 3), the log-likelihood values used in the reward varied significantly in magnitude, destabilizing learning significantly, leading to policies that could no longer be numerically evaluated. This result means that stationarity is a critical component of CSIL.

N.9 Effectiveness of constant rewards

Many imitation rewards typically encode an inductive bias of positivity or negativity, which encourage survival or 'minimum time' strategies respectively [40]. We ablate CSIL with this inductive bias for survival tasks ( Hopper-v2 and Walker-v2 ) and minimum-time tasks ( robomimic ). Figure 47 shows that positive-only rewards are sufficient for Hopper-v2 but not for Walker-v2 . Figure 48 shows that negative-only rewards is effective for the simpler tasks ( Lift ) but not the harder tasks ( NutAssemblySquare) .

N.10 Importance of behavioral cloning pre-training

Figure 49 shows CSIL without any policy pre-training, so the reward is learned only though the finetuning objective. CSIL without pre-training does not appear to work for any of the tasks. This poor performance could be attributed to the lack of hyperparameter tuning, since the reward learning hyperparameters were chosen assuming a good initialization. It could also due to the maximum likelihood fitting of the coherent reward, since the refinement does not use the 'faithful' objective (Section D) that improves the regression fit.

0.0

0.5

1.0

Return

(

n

=1)

HalfCheetah-v2

Hopper-v2

0.0

0.5

1.0

Return

(

n

=3)

0.0

0.5

1.0

Return

(

n

=10)

0

0.25

0.5

0.75

1.0

Steps (

1×10 6

)

0.0

0.5

1.0

Return

(

n

=30)

csil (w/o critic regularization)

csil

0

0.25

0.5

0.75

1.0

Steps (

1×10 6

)

0

0.25

0.5

0.75

1.0

Steps (

1×10 6

)

0

0.25

0.5

0.75

1.0

Steps (

1×10 6

)

0.0

0.5

1.0

Return

(

n

=1)

Hopper-v2

0.0

0.5

1.0

Return

(

n

=3)

0.0

0.5

1.0

Return

(

n

=10)

0

0.25

0.5

0.75

1.0

Steps (

1×10 6

)

0.0

0.5

1.0

Return

(

n

=30)

0

0.25

0.5

0.75

1.0

Steps (

1×10 6

)

csil (w/ reward = 1)

csil

0.0

0.5

1.0

Success

(

n

=200)

Lift

PickPlaceCan

NutAssemblySquare

0.0

0.5

1.0

Success

(

n

=100)

0.0

0.5

1.0

Success

(

n

=50)

0

0.25

0.5

0.75

1.0

Steps (

1×10 6

)

0.0

0.5

1.0

Success

(

n

=25)

0

0.25

0.5

0.75

1.0

Steps (

1×10 6

)

csil (-1 reward)

csil

0

0.25

0.5

0.75

1.0

Steps (

1×10 6

)

0.0

0.5

1.0

Return

(

n

=1)

HalfCheetah-v2

Hopper-v2

0.0

0.5

1.0

Return

(

n

=3)

0.0

0.5

1.0

Return

(

n

=10)

0

0.25

0.5

0.75

1.0

Steps (

1×10 6

)

0.0

0.5

1.0

Return

(

n

=30)

csil (w/o BC pretraining)

csil

0

0.25

0.5

0.75

1.0

Steps (

1×10 6

)

0

0.25

0.5

0.75

1.0

Steps (

1×10 6

)

0

0.25

0.5

0.75

1.0

Steps (

1×10 6

)

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper studies the imitation learning (IL) problem. There are two main classes of IL methods: behavioral cloning (BC) and inverse reinforcement learning (IRL), each with its own advantages. This paper proposes an IL method Coherent Soft Imitation Learning (CSIL) which combines BC and IRL. CSIL first learns a reward function named “coherent reward” under which the BC policy is optimal, and then performs soft policy iteration on the learned coherent reward. The authors evaluate CSIL on a wide range of settings and tasks including online and offline IL for high-dimensional and image-based continuous control tasks.\n\nSTRENGTHS:\n1. The proposed algorithm CSIL is novel according to my knowledge.\n2. The authors evaluate CSIL on a wide range of settings and tasks including tabular tasks, online and offline IL from agent demonstrations, IL from human demonstrations from states and images.\n\nWEAKNESSES:\n1. The proposed method CSIL fails to combine the advantages of both BC and IRL, which is the main contribution claimed in this paper. In CSIL, the coherent reward function is learned under the principle that the BC policy is the (soft) optimal policy with that reward function. So the best policy we can expect to obtain from performing soft policy iteration on the learned coherent reward function is the BC policy. Besides, CSIL does not retain the main merit of IRL: addressing the compounding errors issue of BC. It is known that IRL leverages the state-action distribution matching principle to address the compounding errors issue in BC [1, 2, 3]. However, CSIL does not inherit this principle to learn the reward function. Therefore, the advantage of CSIL over BC and IRL is unclear.\n2. The theory about the coherent reward (Theorem 1) is not novel. In [4], they have derived the feasible reward formula in the maximum entropy RL setting, which is exactly the case when the prior policy is a uniform policy in Theorem 1 (this choice of prior policy is also used in experiments in this paper). However, the authors do not discuss this existing result.\n3. The writing in this paper needs improvement, and some parts are confusing. Even though I have read most of the IL papers in the reference, I still found it difficult to follow this paper. Please refer to the detailed comments in the Question section.\n4. In experiments, CSIL performs worse than the existing method DAC on the online IL in MuJoCo tasks. DAC outperforms CSIL in HalfCheetah-v2 and Hopper-v2 while having a competitive performance with CSIL in Ant-v2 and Walker2d-v2. Besides, the offline IL setting considered in this paper is actually the offline IL with a supplementary dataset setting instead of the pure offline IL setting. However, the authors does include the methods in the offline IL with a supplementary dataset setting [5, 6] for comparison.\n\nReferences:\n\n[1] Seyed Kamyar Seyed Ghasemipour, Richard Zemel, and Shixiang Gu. A divergence minimization perspective on imitation learning methods. In Conference on Robot Learning, 2019.\n\n[2] Nived Rajaraman, Lin Yang, Jiantao Jiao, and Kannan Ramchandran. Toward the fundamental limits of imitation learning. Advances in Neural Information Processing Systems, 2020.\n\n[3] Tian Xu, Ziniu Li, and Yang Yu. Error bounds of imitating policies and environments. Advances in Neural Information Processing Systems, 2020.\n\n[4] Haoyang Cao, Samuel Cohen, and Lukasz Szpruch. Identifiability in inverse reinforcement learning. In Advances in Neural Information Processing Systems, 2021.\n\n[5] Geon-Hyeong Kim, Seokin Seo, Jongmin Lee, Wonseok Jeon, HyeongJoo Hwang, Hongseok Yang, and Kee-Eung Kim. DemoDICE: Offline imitation learning with supplementary imperfect demonstrations. In Proceedings of the 10th International Conference on Learning Representations, 2022.\n\n[6] Haoran Xu, Xianyuan Zhan, Honglei Yin, and Huiling Qin. Discriminator-weighted offline imitation learning from suboptimal demonstrations. In Prooceedings of the 39th International Conference on Machine Learning, 2022.\n\nQUESTIONS:\n1. Was Figure 1 obtained from experimental results? If so, please provide a brief description of the experimental setup used to obtain the results shown in Figure 1.\n2. Algorithm 1 does not provide a clear description of the proposed CSIL algorithm. First, it does not explain how to initialize the shaped critic and only shows a fixed point equation where both sides depend on the critic to be learned. Additionally, Algorithm 1 does not clearly demonstrate how CSIL utilizes online interactions, offline samples, or the dynamic model.\n\nLIMITATIONS:\nPlease see the detailed reviews in the Weakness part.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-06-26"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis study seeks to leverage the advantages of Behavioral Cloning (BC) and Inverse Reinforcement Learning (IRL) to develop a sample-efficient imitation learning algorithm. However, the integration of these two approaches is not straightforward, as optimizing the policy using a dynamic reward diminishes the benefits of BC pre-training. To address this challenge, this work derives a shaped reward based on pre-trained BC according to the entropy-regularized policy update. This shaped reward enables policy refinement without compromising the advantages gained from BC, as it remains consistent with BC. Experimental results demonstrate that the proposed Coherent Soft Imitation Learning (CSIL) effectively addresses both online and offline imitation learning tasks in high-dimensional continuous control and image-based scenarios.\n\nSTRENGTHS:\n* The experiments are extensive and comprehensive.\n* The proposed algorithm is able to solve complex control tasks using only a few demonstrations.\n* The novelty is commendable.\n\nWEAKNESSES:\n* The performance is influenced a lot by the efficacy of BC.\n* The algorithm requires some tricks to work well.\n* The readability is not good.\n\nQUESTIONS:\n* How is p(a|s) in Eq 2 derived? It doesn’t seem to appear in [1].\n* How is p(a|s) in Eq 4 derived? It doesn’t seem to appear in [2].\n* Definition 1 appears to be somewhat redundant, given that the information is already provided in Eq 4 and KL-regularized RL objective in the Background section.\n* What does \"classifier\" refer to in Fig 1? Does it correspond to the discriminator in GAIL?\n* In Fig 1, it appears that the color intensity of Ours does not necessarily indicate a greater distance from the data. Does this imply that the agent may be unable to return to expert support in certain regions?\n* Algorithm 1 is vague, for example:\n    * Where is the fine tuning part using additional data(online/offline).\n    * The shape coherent reward uses $\\theta$ but there is no $\\theta$ at the right hand side of $=$.\n    * What’s the exact objective/equation for computing $\\tilde{Q}_n$.\n    * It would be better to include the tricks mentioned in Section 4 into Algorithm 1.\n* The placement of “|” in Lemma 2 seems to be wrong: $p(Y, w| X)$ -> $p(Y | w, X)$\n* What is the influence of the reference policy ($p(a | s)$ or $q_{\\theta_1}(a | s)$)?\n* What would happen if we only use Eq 10 or its improved version to train without BC pre-training?\n* In Fig 27, BC’s performance is better than the initial performance of CSIL, which is different from Fig 31, what’s the reason for this?\n* In Fig 31, there are some tasks where BC’s performance is better than CSIL, e.g HalfCheetah-v2(n=30) and Hopper-v2(n=3). This seems to contradict the feature of CSIL, what’s the reason for this?\n* The scores of the experts are somewhat low at HalfCheetah-v2(8770) and Hopper-v2(2798), it would be better to train on better experts, e.g HalfCheetah-v2($\\ge$12000) and Hopper-v2($\\ge$3500).\n* In Fig 3, why does CSIL exhibit a deteriorating performance in Hopper-v2 as the number of demonstrations increases?\n* Some typos:\n    * Fig 28, again -> against\n    * Fig 36, s hows -> shows\n    * Line 984, and and -> and\n \n1: Brian D Ziebart. Modeling purposeful adaptive behavior with the principle of maximum causal entropy. Carnegie Mellon University, 2010.  \n2: Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: Off- policy maximum entropy deep reinforcement learning with a stochastic actor. In International Conference on Machine Learning, 2018.\n\nLIMITATIONS:\nAs mentioned in the paper, the performance of CSIL is constrained by the capabilities of BC. If BC fails to successfully address the task with an adequate number of demonstrations, CSIL will encounter similar limitations.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-05"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe authors propose a hybrid BC and IRL method, which uses a maximum entropy KL-regularized BC policy to define a shaped reward which can be optimized with IRL.  The authors derive a “coherent” reward which is defined as a reward for which the BC policy is optimal, by inverting the soft policy iteration update.  The full method first trains a BC policy, defines the coherent reward based on the policy, then uses soft policy iteration to optimize this reward.  The authors also provide a version of the method for continuous control which requires several additional components.  They compare against other IRL algorithms and BC in a tabular and continuous environments and in online and offline RL settings.  Overall, the method performs well in continuous environments given enough demonstrations and performs comparably with other methods elsewhere.\n\nSTRENGTHS:\n- The authors apply a novel perspective of coherence to formulating a hybrid BC and IRL method.\n- The proposed method is well-derived with theoretical backing.\n- Extensive connections to prior works throughout.\n- Experiments in multiple different settings and control tasks.\n\nWEAKNESSES:\n- While the tabular form of the method is simple and elegant, the continuous version is quite complex and requires many additional components to work.  This includes training the critic with an additional auxiliary loss, using heteroscedastic regression to fit the expert data, and finetuning the coherent reward with a new objective.  In comparison, most other IL methods can be applied out-of-the-box for both discrete and continuous tasks.  Given its complexity, the proposed method could be difficult to reproduce and tune in new environments.  The authors do demonstrate the continuous method on multiple tasks in different domains and include multiple ablations.  However, many details and all ablations are not included in the main paper.  It would be helpful to at least summarize the results to address the method complexity.\n- The derivation of the method can be difficult to follow and seems overly complicated in places.  For example, my understanding is that the initial policy is trained with KL-regularized BC but the authors use pseudo-likelihoods and stochastic process regression to derive it which seem like interesting connections but not necessary for the derivation.  Providing a simpler derivation or some intuition to guide readers could help make this part clearer.\n\nOverall, these issues are minor, could be addressed with a little re-writing, and do not outweigh the main strengths of the paper.  The authors succeed in deriving and empirically verifying a novel hybrid BC and IRL method with the idea of coherence.  While additional analysis on the robustness of the method and a simplified derivation could help, these are not major concerns that detract from the paper's main contributions.\n\nQUESTIONS:\n- Is $r_\\theta$ updated in each policy iteration with the new policy, $q_{\\theta_i}$?  If so, does this change the coherence property?\n- For continuous control policies, you state that you use the cloned policy as the prior in forming the reward.  Wouldn’t this lead to zero rewards everywhere because $r(s,a) = \\alpha (\\log q_\\theta(s,a) - \\log q_\\theta(s,a))$?\n\nLIMITATIONS:\nYes, the note that the BC policy is not always viable as a coherent reward, especially if the task is not solvable with BC even given many demonstrations.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-06"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper proposes an approach to inverse RL to fine-tune a behavior cloning policy using RL on online of offline data sources. Adopting the KL-regularized view to RL/IRL, CSIL expresses the reward in terms of the behavior cloned policy and connects the result with the well-known reward shaping results. This \"coherent\" reward is thereby used to improve the policy, resulting in a simple and performant inverse RL algorithm, compared to game-theoretic/adversarial approaches.\n\nSTRENGTHS:\nThe paper provides a fairly comprehensive overview and background of inverse RL. The construction of the coherent reward using the entropy-regularized RL framework is interesting, and results in a simple algorithm with soft policy iteration. The experimental results are quite convincing and informative, and sufficiently compared against competitive methods.\n\nWEAKNESSES:\nIn general, the presentation is a bit more jargon heavy than it needs to be, the writing can be greatly simplified to communicate key ideas more clearly, and make them more amenable for uptake.\n\nThe reward refinement seems to invoke a minimax optimization procedure. The paper claims to bypass adversarial IRL methods, but the reward refinement procedure seems to contradict that. Can you clarify the questions below?\n\nIt would be good to clarify where the uniform prior is used and where the reference policy is initialized to the BC cloned policy. More details for how the critic is initialized and trained would help address this question.\n\nQUESTIONS:\nHow important is reward refinement? Do the final experiments use reward refinement? If so, can you report performance with and without reward refinement.\n\nThe paper mentions the issue of reward/critic causing the unlearning of the initialization. How does CSIL avoid this issue?\n\nLIMITATIONS:\nThe paper provides a reasonable discussion of limitations, particularly how the success of CSIL can be sensitive to the success of the initial BC policy.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-08"}