TaSIL: Taylor Series Imitation Learning
Abstract
We propose Taylor Series Imitation Learning (TaSIL), a simple augmentation to standard behavior cloning losses in the context of continuous control.
TaSIL penalizes deviations in the higher-order Taylor series terms between the learned and expert policies.
We show that experts satisfying a notion of incremental input-tostate stability are easy to learn, in the sense that a small TaSIL-augmented imitation loss over expert trajectories guarantees a small imitation loss over trajectories generated by the learned policy.
We provide generalization bounds for TaSIL that scale as ˜ O
(1 /n ) in the realizable setting, for n the number of expert demonstrations.
Finally, we demonstrate experimentally the relationship between the robustness of the expert policy and the order of Taylor expansion required in TaSIL, and compare standard Behavior Cloning, DART, and DAgger with TaSIL-loss-augmented variants.
In all cases, we show significant improvement over baselines across a variety of MuJoCo tasks.
Introduction
Imitation learning (IL), wherein expert demonstrations are used to train a policy [1, 2], has been successfully applied to a wide range of tasks, including self-driving cars [3, 4], robotics [5] and video game playing [6].
While IL is typically more sample-efficient than reinforcement learning-based alternatives, it is also known to be sensitive to distribution shift: small errors in the learned policy can lead to compounding errors and ultimately, system failure [3, 6].
In order to mitigate the effects of distribution shift caused by policy error, two broad approaches have been taken by the community.
On-policy approaches such as DAgger [6] augment the data-set with expert-labeled or corrected trajectories generated by learned policies.
In contrast, off-policy approaches such as DART [7] and
GAIL [8] augment the data-set by perturbing the expert controlled system.
In both cases, the goal is to provide examples to the learned policy of how the expert recovers from errors.
While effective, these methods either require an interactive expert (on-policy), or access to a simulator for policy rollouts during training (off-policy), which may not always be practically feasible.
In this work, we take a more direct approach towards mitigating distribution shift.
Rather than providing examples of how the expert policy recovers from error, we seek to endow a learned policy with the robustness properties of the expert directly.
In particular, we make explicit the underlying assumption in previous work that a good expert is able to recover from perturbations through the notion of incremental input-to-state stability , a well-studied control theoretic notion of robust nonlinear stability.
Under this assumption, we show that if the p -th order Taylor series approximation of the learned and expert policies approximately match on expert-generated trajectories, then the learned policy induces closed-loop behavior similar to that of the expert closed-loop system.
Here the order p is determined by the robustness properties of the expert, and makes quantitative the informal observation that more robust experts are easier to learn.
Fundamentally, we seek to characterize settings where offline imitation learning is Probably Approximately Correctly (PAC)-learnable [9]; that is, defining notions of expert data and designing algorithms such that low training time error on offline expert demonstrations implies low test time error along the learned policy’s trajectories, in spite of distribution shift.
Contributions We propose and analyze Taylor Series Imitation Learning (TaSIL), which augments the standard behavior cloning imitation loss to capture errors between the higher-order terms of the
Taylor series expansion of the learned and expert policies.
We reduce the analysis of TaSIL to the analysis of a supervised learning problem over expert data: in particular, we identify a robustness criterion for the expert such that a small TaSIL-augmented imitation loss over expert trajectories guarantees that the difference between trajectories generated by the learned and expert policies is also small.
We also provide a finite-difference based approximation that is applicable to experts that cannot directly query their higher-order derivatives, expanding the practical applicability of TaSIL.
We show in the realizable setting that our algorithm achieves generalization bounds scaling as ˜ O (1 /n ) , for n the number of expert demonstrations.
Finally, we empirically demonstrate (i) that the relationship between the robustness of the expert policy and the order of Taylor expansion required in TaSIL predicted by our theory is observed in practice, and (ii) the benefits of using the TaSIL-augmented imitation loss by comparing the sample-efficiency of standard and TaSIL-loss-augmented behavior cloning, DART, and DAgger on a variety of MuJoCo tasks: on hard instances where behavior cloning fails, the TaSIL-augmented variants show significant performance gains.
1.1 Related work
Imitation learning Behavior cloning is known to be sensitive to compounding errors induced by small mismatches between the learned and expert policies [3, 6].
On-policy [6] and off-policy [7, 8] approaches exist that seek to prevent this distribution shift by augmenting the data-set created by the expert.
While DAgger is known to enjoy ˜ O ( T ) sample-complexity in the task horizon T for loss functions that are strongly convex in the policy parameters, 3 we are not aware of finite-data guarantees for DART or GAIL.
In addition to these seminal papers, there is a body of work that seeks to leverage control theoretic techniques [10, 11] to ensure (robust) stability of the learned policy.
More closely related to our work are the results by Ren et al.
[12] and Tu et al.
In Ren et al.
[12], a two-stage pipeline of imitation learning followed by policy optimization via a PAC-Bayes regularization term is used to provide generalization bounds of the learned policy across randomly drawn environments.
This work is mostly complementary to ours, as TaSIL could in principle be used to augment the imitation losses used in the first stage of their pipeline (leaving the second stage unmodified).
In Tu et al.
[13], sample-complexity bounds for IL are provided under the assumption that the learned policy can be explicitly constrained to match the incremental stability properties of the expert.
While conceptually appealing, practically enforcing such stability constraints is difficult, and as such Tu et al.
[13] resort to heuristics in their implementation.
In contrast, we provide sample-complexity guarantees for a practically implementable algorithm, and under much milder stability assumptions.
Robust stability and learning for continuous control There is a rich body of work applying
Lyapunov stability or contraction theory [14] to learning for continuous control.
For example [15–18] use stability-based regularizers to trim the hypothesis space, and empirically show that this leads to more sample-efficient and robust learning methods.
Lyapunov stability and contraction theory have also been used to provide finite sample-complexity guarantees for adaptive nonlinear control [19] and learning stability certificates [20].
Problem formulation
We consider the nonlinear discrete-time dynamical system
[Equation]
where x t ∈ R d is the system state, u t ∈ R m is the control input, and f : R d × R m → R d defines the system dynamics.
We study system (1) evolving under the control input u t
=
π
( x t )
+ ∆ t , for π : R d → R m a suitable control policy, and ∆ t ∈ R m an additive input perturbation (that will be used to capture policy errors).
We define the corresponding perturbed closed-loop dynamics by x t +1 = f π cl ( x, ∆ t ) :
= f ( x, π ( x ) + ∆ t ) , and use x πt ( ξ, { ∆ s } t − 1 s =0 ) to denote the value of the state x t at time t evolving under control input u t
= π ( x ) + ∆ t starting from initial condition x 0 = ξ .
To lighten notation, we use x πt ( ξ ) to denote x πt ( ξ, { 0 } t − 1 s =0 ) , and overload ∥·∥ to denote the Euclidean norm for vectors and the operator norm for matrices and tensors.
Initial conditions ξ are assumed to be sampled randomly from a distribution D with support restricted to a compact set X .
We assume access to n rollouts of length T from an expert π ⋆ , generated by drawing initial conditions { ξ i } ni =1 i.i.d.
from D .
The IL task is to learn a policy ˆ π which leads to a closed-loop system with similar behavior to that induced by the expert policy π ⋆ as measured by the expected imitation gap E
ξ max 1 ≤ t ≤ T (cid:13)(cid:13)
x ˆ πt ( ξ ) − x π ⋆ t ( ξ ) (cid:13)(cid:13) .
The baseline approach to IL, typically referred to as behavior cloning (BC), casts the problem as an instance of supervised learning.
Denote the discrepancy between an evaluation policy ¯ π and the expert policy π ⋆ on a trajectory generated by a rollout policy π d starting at initial condition ξ by ∆ π d t ( ξ ; ¯ π ) :
=
[Equation]
−
[Equation]
BC directly solves the supervised empirical risk minimization (ERM) problem
[Equation]
over a suitable policy class Π .
Here, h : ( R m ) T → R is a loss function which encourages the discrepancy terms ∆ π ⋆ t ( ξ i ; π ) to be small along the expert trajectories.
While behavior cloning is conceptually simple, it can perform poorly in practice due to distribution shifts triggered by errors in the learned policy ˆ π bc .
Specifically, due to the effects of compounding errors, the closed-loop system induced by the behavior cloning policy ˆ π bc may lead to a dramatically different distribution over system trajectories than that induced by the expert policy π ⋆ , even when the population risk on the expert data, E ξ h ( { ∆ π ⋆ t ( ξ ; ˆ π bc ) } T − 1 t =0 ) , is small.
As described in the introduction, existing approaches to mitigating distribution shift seek to augment the data with examples of the expert recovering from errors.
These approaches either require an interactive oracle (e.g., DAgger) or access to a simulator for policy rollouts (e.g., DART, GAIL), and may not always be practically applicable.
To address the distribution shift challenge without resorting to data-augmentation, we propose TaSIL, an off-policy IL algorithm which provably leads to learned policies that are robust to distribution shift.
The rest of the paper is organized as follows: in Section 3, we focus on ensuring robustness to policy errors for a single initial condition ξ .
We show that the imitation gap ∥ x πt ( ξ ) − x π ⋆ t ( ξ ) ∥ between a test policy π and the expert policy π ⋆ can be controlled by the TaSIL-augmented imitation loss evaluated on the expert trajectory { x π ⋆ t ( ξ ) } , effectively reducing the analysis of the imitation gap to a supervised learning problem over the expert data.
In Section 4, we integrate these results with tools from statistical learning theory to show that n ≳ ε − r /δ trajectories are sufficient to achieve imitation gap of at most ε with probability at least 1 − δ ; here r > 0 is a constant determined by the stability properties of the expert policy π ⋆ , with more robust experts corresponding to smaller values of r .
Finally, in Section 5, we validate our analysis empirically, and show that using the TaSIL-augmented loss function in IL algorithms leads to significant gains in performance and sample efficiency.
Bounding the imitation gap on a single trajectory
In this section, we fix an initial condition ξ and test policy π , and seek to control the imitation gap Γ T ( ξ ; π ) :
= max 1 ≤ t ≤ T ∥ x πt ( ξ, { 0 } ) − x π ⋆ t ( ξ, { 0 } ) ∥ .
A natural way to compare the closed-loop behavior of a test policy π to that of the expert policy π ⋆ is to view the discrepancy ∆ πt ( ξ ; π ) :
= π ( x πt ( ξ ))
− π ⋆ ( x πt ( ξ )) as an input perturbation to the expert closed-loop system.
By writing x t +1 = f π cl ( x t , 0) =
f π ⋆ cl ( x t , ∆ πt ( ξ ; π )) , the imitation gap can be written as Γ T ( ξ ; π ) =
max t ∥ x π ⋆ t ( ξ, { ∆ πs ( ξ ; π ) } t − 1 s =0 ) − x π ⋆ t ( ξ, { 0 } ) ∥ , suggesting that closed-loop expert systems that are robust to input perturbations, as measured by the difference between nominal and perturbed trajectories, will lead to learned policies that enjoy smaller imitation gaps.
Stability conditions defined in terms of differences between nominal and perturbed trajectories have been extensively studied in robust nonlinear control theory via the notion of incremental input-tostate stability ( δ -ISS)
(see e.g., Angeli [21] and references therein).
Before proceeding, we recall definitions of standard comparison functions [22]: a function γ ( x ) is class K if it is continuous, strictly increasing, and satisfies γ (0) = 0 , and a function β ( x, t ) is class KL if it is continuous, β ( · , t ) is class K for each fixed t , and β ( x, · ) is decreasing for each fixed x .
Definition 3.1 ( δ -ISS system) .
Consider the closed-loop system evolving under policy π and subject to perturbations ∆ t given by x t +1 = f π cl ( x t , ∆ t ) .
The closed-loop system f π cl ( x t , ∆ t ) is incremental input-to-state stable ( δ -ISS) if there exists a class KL function β and a class K function γ such that for all initial conditions ξ 1 , ξ 2 ∈ X , perturbation sequences { ∆ t } t ≥ 0 , and t ∈ N : (cid:13) (cid:13)
[Equation]
Definition 3.1 says that: (i) trajectories generated by δ -ISS systems converge towards each other if they begin from different initial conditions, and (ii) the effect of bounded perturbations { ∆ t } on trajectories is bounded.
Our results will only require the stability conditions of Definition 3.1 to hold for a class of norm bounded perturbations.
In light of this, we say that a system is η -locally δ -ISS if equation (2) holds for all input perturbations satisfying sup t ∈ N ∥ ∆ t ∥ ≤ η .
By writing x t +1 = f π cl ( x t , 0) =
f π ⋆ cl ( x t , ∆ πt ( ξ ; π )) , we conclude that if x t +1 =
f π ∗ cl ( x t , ∆ t ) is η -locally δ -ISS , and that sup t ∈ N ∥ ∆ πt ( ξ ; ˆ π ) ∥ ≤ η , then equation (2) yields
[Equation]
Equation (3) shows that the imitation gap ∥ x πt ( ξ ) − x π ⋆ t ( ξ ) ∥ is controlled by the maximum discrepancy ∥ ∆ πt ( ξ ; π ) ∥ incurred on trajectories generated by the policy π .
A natural way of bounding the test discrepancy ∥ ∆ πt ( ξ ; π ) ∥ , defined over trajectories generated by π , by the training discrepancy ∥ ∆ π ⋆ t ( ξ ; π ) ∥ , defined over trajectories generated by π ⋆ , is to write out a Taylor series expansion of the former around the latter, i.e., to write:
[Equation]
where ∂ x ∆ π ⋆ t ( ξ ; π ) denotes the partial derivative of the discrepancy with respect to the argument of the policies π and π ⋆ .
The challenge however, is that the imitation gap Γ T ( ξ ; π ) also appears in the
Taylor series expansion (4), leading to an implicit constraint.
We show that this can be overcome if the order p of the Taylor series expansion (4) is sufficiently large, as determined by the decay rate of the class K function γ ( · ) defining the robustness of the expert policy.
We begin by identifying a condition reminiscent of adversarially robust training objectives that ensures a small imitation gap.
Proposition 3.1.
Let the expert closed-loop system f π ⋆ cl be η -locally δ -ISS for some η > 0 .
Fix an imitation gap bound ε > 0 , initial condition ξ , and policy π .
Then if
[Equation]
we have that the imitation gap satisfies Γ T ( ξ ; π ) ≤ ε .
Proposition 3.1 states that if a policy π is sufficiently close to the expert policy π ⋆ in a tube around the expert trajectory, then the imitation gap remains small.
How to ensure that inequality (5) holds using only offline data from expert trajectories is not immediately obvious.
The Taylor series expansion (4) suggests that a natural approach to satisfying this condition is to match derivatives of the test policy π to those of the expert policy π ⋆ .
We show next that a sufficient order for such a Taylor series expansion is naturally determined by the decay rate of the class K function γ ( x ) towards 0 .
Less robust experts have functions that decay to 0 more slowly and will lead to stricter sufficient conditions.
Conversely, more robust experts have functions that decay to 0 more quickly, and will lead to more relaxed sufficient conditions.
We focus on two disjoint classes of class K functions: (i) functions that decay in their argument faster than a linear function, i.e. γ ( x ) <
O ( x ) as x → 0 + , and (ii) functions that decay in their argument no faster than a linear function, i.e. γ ( x ) ≥ Ω( x ) as x → 0 + .
Rapidly decaying class K functions
We show that when the class K function γ ( x ) decays to 0 faster than O ( x ) in some neighborhood of 0 , then matching the zeroth-order difference max t ∥ ∆ π ⋆ t ( ξ ; π ) ∥ on the expert trajectory, as is done in vanilla behavior cloning, suffices to close the imitation gap.
We make the following assumption on the test policy π and expert policy π ⋆ .
Assumption 3.1.
There exists a non-negative constant L π such that ∥ ¯ π ( x ) − ¯ π ( y ) ∥ 2 ≤ L π ∥ x − y ∥for all x , y , and ¯ π ∈ { π, π ⋆ } .
Proposition 3.1 then leads to the following guarantee on the imitation gap.
Theorem 3.1.
Fix a test policy π and initial condition ξ ∈ X , and let Assumption 3.1 hold.
Let f π ⋆ cl be η -locally δ -ISS for some η > 0 , and assume that the class K function γ ( · ) in (2) satisfies γ ( x ) ≤ O ( x 1+ r ) for some r > 0 .
Choose constants µ, α > 0 such that
[Equation]
for all 0 ≤ x ≤ α .
Provided that the imitation error on the expert trajectory incurred by π satisfies:
[Equation]
then for all 1 ≤ t ≤ T the instantaneous imitation gap is bounded as
[Equation]
Theorem 3.1 shows that if a policy π is a sufficiently good approximation of the expert policy π ⋆ on an expert trajectory { x π ⋆ t ( ξ ) } , then the imitation gap ∥ x π ⋆ t ( ξ ) − x πt ( ξ ) ∥ can be upper bounded in terms of the discrepancy term max 0 ≤ k ≤ t − 1 ∥ ∆ π ⋆ t ( ξ ; π ) ∥ evaluated on the expert trajectory .
To help illustrate the effect of the decay parameter r on the choices of µ and α , we make condition (6) more explicit by assuming that γ ( x ) ≤ Cx 1+ r for all x ∈ [0 , 1] .
Then one can choose µ = 2 1+ r C and α
= O (1)( L 1+ r π C ) − 1 /r .
This makes clear that a larger r , i.e., a more robust expert, leads to less restrictive conditions (7) on the policy π and a tighter upper bound on the imitation gap (8) (assuming max 0 ≤ k ≤ t − 1 ∥ ∆ π ⋆ t ( ξ ; π ) ∥ < 1 ).
In particular, for such systems, vanilla behavior cloning is sufficient to ensure bounded imitation gap.
This also makes clear how the result breaks down when r ≈ 0 , i.e., when the decay rate is nearly linear, as the neighborhood α can become arbitrarily small, such that it may be impossible to learn a policy π that satisfies the bound (7) with a practical number of samples n .
The interplay of the bounds (7) in Theorem 3.1 (and the subsequent theorems of its like) and the sample-complexity of imitation learning will be discussed in further detail in Section 4.
Slowly decaying class K functions
When the class K function γ ( x ) decays to 0 slowly, for reasons discussed above, controlling the zeroth-order difference ∆ π ⋆ t ( ξ ; π ) may not be sufficient to bound the imitation gap.
In particular, we consider class K functions satisfying γ ( x ) ≤ O ( x 1 /r ) for some r ≥ 1 .
Setting p =
⌊ r ⌋ , we now show that matching up to the p -th total derivative of π ⋆ is sufficient to control the imitation gap.
Analogously to Assumption 3.1, we make the following regularity assumption on the test policy π and expert policy π ⋆ .
Assumption 3.2.
For a given non-negative p ∈ N , assume that the test policy π and expert policy π ⋆ are p -times continuously differentiable, and there exists a constant L ∂ p π ≥
such that
[Equation]
for all x, x 0 and ¯ π ∈ { π, π ⋆ } , where (cid:0) J px 0 ¯ π (cid:1) ( x ) := (cid:80) pj =
1 j ! (cid:0) ∂
jx ¯ π ( x 0 ) (cid:1) ( x − x 0 ) ⊗ j is the p -th order
Taylor polynomial of ¯ π evaluated at x 0 , and ⊗ denotes the tensor product.
With this assumption in hand
[Equation]
we provide the following guarantee on the imitation gap.
Theorem 3.2.
Let f π ⋆ cl be η -locally δ -ISS for some η > 0 , and assume that the class K function γ ( · ) in (2) satisfies γ ( x ) ≤ O ( x 1 /r ) for some r ≥ 1 .
Fix a test policy π and initial condition ξ ∈ X
[Equation]
and let Assumption 3.2 hold for p ∈ N satisfying p + 1 − r > 0 .
Choose µ, α > 0 such that
[Equation]
for
[Equation]
Provided the j th total derivatives, j = 0 , . . .
, p , of the imitation error on the expert trajectory incurred by π satisfy:
[Equation]
then for all 1 ≤ t ≤ T the instantaneous imitation gap is bounded by
[Equation]
Theorem 3.2 shows that if the p -th order Taylor series of the policy π approximately matches that of the expert policy π ⋆ when evaluated on an expert trajectory { x π ⋆ t ( ξ ) } , then the imitation gap ∥ x π ⋆ t ( ξ ) − x πt ( ξ ) ∥ can be upper bounded in terms of the derivatives of the discrepancy term, i.e., by max 0 ≤ k ≤ t − 1 max 0 ≤ j ≤ p (cid:13)(cid:13)
∂ jx ∆ π ⋆ k ( ξ ; π ) (cid:13)(cid:13) , evaluated on the expert trajectory .
To help illustrate the effect of the choice of the order p on the constants µ and α , we make condition (10) more explicit by assuming that γ ( x ) ≤ Cx 1 /r for all x ∈ [0 , 1] .
Then one can choose µ = 2 1 /r C and α
= O (1)( L ∂ p π C r ) − 1 / ( p +1 − r ) .
This expression highlights a tradeoff: by picking larger order p , the right hand side α of bound (11) increases, but at the expense of having to match higher-order derivatives.
This also highlights that both the order p and closeness required by Equation (11) get increasingly restrictive as r increases, matching our intuition that less robust experts lead to harder imitation learning problems.
Using estimated derivatives We show in Appendix C that the results of Theorems 3.1 and 3.
extend gracefully to when only approximate derivatives (cid:91) ∂
jx π
⋆ ( x ) can be obtained, e.g., through finite-difference methods.
In particular, if ∥ (cid:91) ∂ jx π ⋆ ( x ) − ∂ jx π ⋆ ( x ) ∥ ≤ ε , then it suffices to appropriately tighten the constraints (11) and (12) by O ( ε 1 /r ) and O ( ε ) , respectively.
Please refer to Appendix C for more details.
Algorithms and generalization bounds for TaSIL
The analysis of Section 3 focused on a single test policy π and initial condition ξ .
Theorems 3.1 and 3.2 motivate defining the p -TaSIL loss function:
[Equation]
The corresponding policy ˆ π TaSIL ,p is the solution to the empirical risk minimization (ERM) problem:
[Equation]
which explicitly seeks to learn a policy π ∈ Π , for Π a suitable policy class, that matches the p -th order Taylor series expansion of the expert policy.
In this section, we analyze the generalization and sample-complexity properties of the p
-TaSIL
ERM problem (15).
Our analysis in this section focuses on the realizable setting: we assume that for every dataset of expert trajectories {{ x π ⋆ t ( ξ i ) } T − 1 t =0 } ni =1 , there exists a policy π ∈ Π that achieves (near) zero empirical risk.
This is true if, for example, π ⋆ ∈ Π .
In this setting, we demonstrate that we can attain generalization bounds that decay as ˜ O ( n − 1 ) , where ˜ O ( · ) hides poly-log dependencies on n .
These rates are referred to as fast rates in statistical learning, since they decay faster than the n − 1 /rate prescribed by the central limit theorem.
We present analysis for the non-realizable setting in
Appendix B: this analysis is standard and yields generalization bounds scaling as O ( n − 1 / 2 ) .
Let G ⊂ [0 , 1]
X be a set of functions mapping some domain X to [0 , 1] .
Let D be a distribution with support restricted to X , and denote the mean of g ∈ G with respect to x ∼ D by E x
[ g ] .
Similarly, fixing data points x 1 , . . . , x n ∈ X , we denote the empirical mean of g by E n [ g ] := n − 1 (cid:80) ni
=1 g ( x i ) .
We focus our analysis on the following class of parametric Lipschitz function classes.
Definition 4.1 (Lipschitz parametric function class) .
A parametric function class G ⊂ [0 , 1] X is called ( B θ , L θ , q ) - Lipschitz if G = { g θ ( · ) | θ ∈ Θ } with Θ ⊂ R q , and it satisfies the following boundedness and uniform Lipschitz conditions:
[Equation]
We assume without loss of generality that B θ L θ ≥ 1 .
The description (16) is very general, and as we show next, is compatible with feed-forward neural networks with differentiable activation functions.
We then have the following generalization bound, which adapts [23, Corollary 3.7] to Lipschitz parametric function classes using the machinery of local
Rademacher complexities [23].
Alternatively, the result can also be derived from [24, Theorem 3].
Theorem 4.1.
Let G ⊂ [0 , 1]
X be a ( B θ , L θ , q ) -Lipschitz parametric function class.
There exists a universal positive constant K < 10 6 such that the following holds.
Given δ ∈ (0 , 1) , with probability at least 1 − δ over the i.i.d.
draws x 1 , . . .
, x n ∼ D , for all g ∈ G , the following bound holds:
[Equation]
We now use Theorem 4.1 to analyze the generalization properties of the p
-TaSIL ERM problem (15).
In what follows, we assume that the expert-closed loop system is stable in the sense of Lyapunov, i.e., that there exists B X > 0 such that sup t ∈ N ,ξ ∈X ∥ x π ⋆ t ( ξ ) ∥ ≤ B X , and consider the following parametric class of p + 2 continuously differentiable policies:
[Equation]
[Equation]
Define the constants
[Equation]
for j = 0 , . . .
, p , and note that they are guaranteed to be finite under our regularity assumptions.
Finally, define the loss function class:
[Equation]
From a repeated application of Taylor’s theorem, we show in Lemma B.2 that B − 1 ℓ,p ( ℓ π ⋆ p ◦
Π θ,p ) is a ( B θ , B − 1 ℓ,p L ℓ,p , q ) -Lipschitz parametric function class for B ℓ,p :
= 2 p +1 (cid:80) pj =
B j and L ℓ,p :=
B X p +1 (cid:80) pj =
L j .
We now combine this with Theorem 4.1 to bound the population risk achieved by the solution to the TaSIL ERM problem (15).
Corollary 4.1.
Let the policy class Π θ,p be defined as in (18) , and assume that π ⋆ ∈ Π θ,p .
Let the function class ℓ π ⋆ p ◦
Π θ,p be defined as in (19) , and constants B ℓ,p , L ℓ,p be defined as above.
Let ˆ π TaSIL ,p be any empirical risk minimizer (15) .
Then with probability at least 1 − δ over the initial conditions { ξ i } n i =1 i .
i .
d . ∼ D n ,
[Equation]
We note that since these generalization bounds solely concern the supervised learning problem of matching the expert on the expert trajectories, the constants do not depend on the stability properties of the trajectories generated by the learned policy ˆ π TaSIL ,p .
To convert the generalization bound (20) to a probabilistic bound on the imitation gap Γ T ( ξ ; ˆ π TaSIL ,p ) , we first apply Markov’s inequality to bound the probability that the conditions of Theorem 3.1 or 3.2 hold by the expected TaSIL loss (20), and then apply Corollary 4.1 together with Markov’s inequality.
Theorem 4.2 (Rapidly decaying class K functions) .
Assume that π ⋆ ∈
Π θ, 0 and let the assumptions of Theorem 3.1 hold for all π ∈ Π θ, 0 .
Let Equation (6) hold with constants µ, α > 0 , and assume without loss of generality that α/µ ≤ 1 , L π µ ≥ 1 / 2 .
Let ˆ π TaSIL , 0 be an empirical risk minimizer of ℓ π ⋆ 0 over the policy class Π θ, 0 for initial conditions { ξ i } i .
i .
d . ∼ D n .
Fix a failure probability δ ∈ (0 , 1) , and assume that
[Equation]
where κ α := q ( µ/α ) 1 / (1+ r ) , κ η := qL π µ/η .
Then with probability at least 1 − δ , the imitation gap evaluated on ξ ∼ D (drawn independently from { ξ i } ni =1 ) satisfies
[Equation]
n
[Equation]
Theorem 4.3 (Slowly decaying class K functions) .
Assume that π ⋆ ∈ Π θ,p , and let the assumptions of Theorem 3.2 hold for all π ∈ Π θ,p .
Let Equation (10) hold with constants µ, α > 0 .
Let ˆ π TaSIL ,p be an empirical risk minimizer of ℓ π ⋆ p over the policy class Π θ,p for initial conditions { ξ i } i .
i .
d . ∼ D n .
Fix a failure probability δ ∈ (0 , 1) , and assume
[Equation]
where κ α,j := (cid:0) µα (cid:1) r pqj !
and κ η,j := (cid:0) L ∂pπ ( p +1)!
µ p + ( j !)
( p +1) /r
+ 1 j ! (cid:1) pqηδ .
Then with probability at least 1 − δ , the imitation gap evaluated on ξ ∼ D (drawn independently from { ξ i } ni =1 ) satisfies
[Equation]
In the rapidly decaying setting, corresponding to more robust experts, Theorem 4.2 states that n ≳ ε − 11+ r /δ expert trajectories are sufficient to ensure that the imitation gap Γ T ( ξ ; ˆ π TaSIL , 0 ) ≲ ε with probability at least 1 − δ .
Recall that more robust experts have larger values of r > 0 , leading to smaller sample-complexity bounds.
In contrast, to achieve the same guarantees on the imitation gap in the slowly decaying setting, Theorem 4.3 states n ≳ ε − r /δ expert trajectories are sufficient, where we recall that less robust experts have larger values of r ≥ 1 .
These theorems quantitatively show how the robustness of an underlying expert affects the sample-complexity of IL, with more robust experts enjoying better dependence on ε than less robust experts.
We note that analogous dependencies on expert stability are reflected in the burn-in requirements, i.e., the number of expert trajectories required to ensure with high probability no catastrophic distribution shift occurs, of each theorem.
Experiments
We compare three standard imitation learning algorithms, Behavior Cloning, DAgger, and DART, to
TaSIL-augmented loss versions.
In TaSIL-augmented algorithms, we replace the standard imitation loss function
[Equation]
with the p -TaSIL-augmented loss
[Equation]
M ea n
P o li c y D
i s c
r e p a n c y
TaSIL order TaSIL order
TaSIL order0.
I m it a ti on
G a p
TaSIL order TaSIL order
TaSIL order where { λ j } pj =0 are positive tunable regularization parameters.
We use the Euclidean norm of the vectorized error in the derivative tensors as more optimizer-amenable surrogate to the operator norm.
We experimentally demonstrate (i) the effect of the expert stability properties and order of the TaSIL loss, and (ii) that TaSIL-loss-based imitation learning algorithms are more sample efficient than their standard counterparts.
All experiments 7 are carried out using Jax [25] GPU acceleration and automatic differentiation capabilities to compute the higher-order derivatives, and the Flax [26] neural network and Optax [27] optimization toolkits.
Stability Experiments To illustrate the effect of the expert closed-loop system stability on samplecomplexity, we consider a simple δ -ISS stable dynamical system with a tunable γ input-to-state gain function.
For state and input x t , u t ∈ R 10 , the dynamics are:
[Equation]
The perturbation function h : R 10 → R 10 is set to a randomly initialized MLP with two hidden layers of width 32 and GELU [28] activations such that the expert π ⋆ ( x ) = − h
( x ) yield a closed loop system f π ⋆ cl ( x, ∆) which is δ -ISS stable with the specified class K function γ (see Appendix D).
We use η = 0 .
95 for all experiments presented here.
We investigate the performance of p -TaSIL loss functions for δ -ISS system with different class K stability.
We sweep K functions γ ( x ) =
Cx ν for ν ∈ [0 . 05 , 3] , C = 5 and p -TaSIL loss functions for p ∈ { 0 , 1 , 2 } (additional details can be found in Appendix D).
The results of this sweep are shown in Figure 1.
Higher-order p -TaSIL losses significantly reduce both the imitation gap and the mean policy discrepancy on test trajectories.
Notably, the first and second order TaSIL loss maintain their improved performance for slower decaying class K functions.
Theorem 3.1 and Theorem 3.2 yield lower bounds of ν = 1 , ν = 2 − 1 , and ν = 3 − 1 for closing the imitation gap using the 0 -TaSIL, 1 -TaSIL, and 2 -TaSIL losses respectively.
Figure 1 demonstrates significant performance degradation in policy discrepancy and decaying imitation gap starting around these threshold values.
MuJoCo Experiments We evaluate the ability of the TaSIL loss to improve performance on standard imitation learning tasks by modifying Behavior Cloning, DAgger [6], and DART [7] to use the ℓ TaSIL , 1 loss and testing them in simulation on different OpenAI Gym MuJoCo tasks [29].
The
MuJoCo environments we use and their corresponding (state, input) dimensions are: Walker2d-v ( 17 , 6 ), HalfCheetah-v3 ( 17 , 6 ), Humanoid-v3 ( 376 , 17 ), and Ant-v3 ( 111 , 8 ).
For all environments we use pretrained expert policies obtained using Soft Actor Critic reinforcement learning by the Stable-Baselines3 [30] project.
The experts consist of Multi-Layer Perceptrons with two hidden layers of 256 units each and ReLU activations.
For all environments, learned policies have 2 hidden layers with 512 units each and GELU activations in addition to Batch Normalization.
The final policy output for both the expert and learned policy are rescaled to the valid action space
[Figure]
w a r d after applying a tanh nonlinearity.
We used trajectories of length T = 300 for all experiments.
We refer to Appendix E for additional experiment details.
In Figure 2 we report the mean expert-normalized rewards across 5 seeds for all algorithms and environments as a function of the trajectory budget provided.
Algorithms with 1-TaSIL loss showed significant improvement in sample-complexity across all challenging environments.
The expert for the Swimmer environment is very robust due to the simplicity of the task, and so as predicted by
Theorem 3.1 and 4.2, all algorithms (with the exception of the vanilla DAgger algorithm due to it initially selecting poor rollout trajectories) are able to achieve near expert performance across all trajectory budgets.
We are also able to nearly match or exceed the performance of standard on-policy methods DAgger and DART with our off-policy 1-TaSIL loss Behavior Cloning in all environments.
Additional experimental results can be found in the appendix.
These include representative videos of the behavior achieved by the expert, BC, and TaSIL-augmented BC policies in the supplementary material, where once again, a striking improvement is observed, especially in low-data regimes for harder environments, as well as a systematic study of finite-difference-based approximations of 1 -TaSIL which achieve comparable performance to Jacobian-based implementations.
Conclusion
We presented Taylor Series Imitation Learning (TaSIL), a simple augmentation to behavior cloning that penalizes deviations in the higher-order Tayler series terms between the learned and expert policies.
We showed that δ -ISS experts are easier to learn, both in terms of the loss-function that needs to be optimized and sample-complexity guarantees.
Finally, we showed the benefit of using
TaSIL-augmented losses in BC, DAgger, and DART across a variety of MuJoCo tasks.
This work opens up many exciting future directions, including extending TaSIL to pixel-based IL, and to (offline/inverse) reinforcement learning settings.
Acknowledgements
We thank Vikas Sindhwani, Sumeet Singh, Jean-Jacques E. Slotine, Jake Varley, and Fengjun Yang for helpful feedback.
Nikolai Matni is supported by NSF awards CPS-2038873, CAREER award
[-Let B ℓ,p := 2 p +1 (cid:80) pj =
-]
[-= B X p +1 (cid:80) pj =
-]
[-=
-]
[-p +1 (cid:80) pj =
-]
[-ni =
-]
[-= 
-]
[-R p +
-]
ECCS-2045834, and a Google Research Scholar award.