Title: Provable Guarantees for Generative Behavior Cloning: Bridging Low-Level Stability and High-Level Behavior
Venue: NeurIPS 2023
Accepted: yes
Keywords: Imitation Learning

Control

Diffusion Models

Optimal Transport
Abstract: We propose a theoretical framework for studying behavior cloning of complex expert demonstrations using generative modeling.
Our framework invokes low-level controllers - either learned or implicit in position-command control - to stabilize imitation around expert demonstrations. We show that with (a) a suitable low-level stability guarantee and (b) a powerful enough generative model as our imitation learner,  pure supervised behavior cloning can generate trajectories matching the per-time step distribution of essentially arbitrary expert trajectories in an optimal transport cost. Our analysis relies on a stochastic continuity property of the learned policy we call "total variation continuity" (TVC). We then show that TVC can be ensured with minimal degradation of accuracy by combining a popular data-augmentation regimen with a novel algorithmic trick: adding augmentation noise at execution time. We instantiate our guarantees for policies parameterized by diffusion models and prove that if the learner accurately estimates the score of the (noise-augmented) expert policy, then the distribution of imitator trajectories is close to the demonstrator distribution in a natural optimal transport distance. Our analysis constructs intricate couplings between noise-augmented trajectories, a technique that may be of independent interest. We conclude by empirically validating our algorithmic recommendations, and discussing implications for future research directions for better behavior cloning with generative modeling.

Provable Guarantees for Generative Behavior Cloning: Bridging Low-Level Stability and High-Level Behavior

Abstract

We propose a theoretical framework for studying behavior cloning of complex expert demonstrations using generative modeling. Our framework invokes low-level controllers - either learned or implicit in position-command control - to stabilize imitation around expert demonstrations. We show that with (a) a suitable low-level stability guarantee and (b) a powerful enough generative model as our imitation learner, pure supervised behavior cloning can generate trajectories matching the per-time step distribution of essentially arbitrary expert trajectories in an optimal transport cost. Our analysis relies on a stochastic continuity property of the learned policy we call 'total variation continuity" (TVC). We then show that TVC can be ensured with minimal degradation of accuracy by combining a popular data-augmentation regimen with a novel algorithmic trick: adding augmentation noise at execution time. We instantiate our guarantees for policies parameterized by diffusion models and prove that if the learner accurately estimates the score of the (noise-augmented) expert policy, then the distribution of imitator trajectories is close to the demonstrator distribution in a natural optimal transport distance. Our analysis constructs intricate couplings between noise-augmented trajectories, a technique that may be of independent interest. We conclude by empirically validating our algorithmic recommendations, and discussing implications for future research directions for better behavior cloning with generative modeling.

1 Introduction

Training dynamic agents from datasets of expert examples, known as imitation learning , promises to take advantage of the plentiful demonstrations available in the modern data environment, in an analogous manner to the recent successes of language models conducting unsupervised learning on enormous corpora of text [67, 70]. Imitation learning is especially exciting in robotics, where mass stores of pre-recorded demonstrations on Youtube [1] or cheaply collected simulated trajectories [42, 20] can be converted into learned robotic policies.

For imitation learning to be a viable path toward generalist robotic behavior, it needs to be able to both represent and execute the complex behaviors exhibited in the demonstrated data. An approach that has shown tremendous promise is generative behavior cloning: fitting generative models, such

Cambridge, MA

x

1

x

2

x

3

x

4

x

5

Smoothing Augmentations

Right Mode

Primitive Controller

Stabilization

as diffusion models [2, 19, 33], to expert demonstrations with pure supervised learning. In this paper, we ask: When can generative behavior cloning imitate arbitrarily complex expert behavior?

In this paper, we are interested in how algorithmic choices interface with the dynamics of the agent's environment to render imitation possible. The key challenge separating imitation learning from vanilla supervised learning is one of compounding error : when the learner executes the trained behavior in its environment, small mistakes can accumulate into larger ones; this in turn may bring the agent to regions of state space not seen during training, leading to larger-still deviations from intended trajectories. Without the strong requirement that the learner can interactively query the expert at new states [40, 57], it is well understood that ensuring some form of stability in the imitation learning procedure is indispensable [68, 27, 50]. While many natural notions of stability exist for simple behaviors, how to enforce stability when imitating more complex behaviors remains an open question. Multi-modal trajectories present a key example of this challenge: consider a robot navigating around an obstacle; because there is no difference between navigating around the object to the right and around to the left, the dataset of expert trajectories may include examples of both options. This bifurcation of good trajectories can make it difficult for the agent to effectively choose which direction to go, possibly even causing the robot to oscillate between directions and run into the object. [19]. Moreover, human demonstrators correlate current actions with the past in order to commit to either a right or left path, which makes even formulating the idea of an 'expert policy ' a conceptually challenging one. Lastly, bifurcations are necessarily incompatible with previous notions of stability derived from classical control theory [68, 27, 50]. In this work, we investigate how these strong and often unrealistic assumptions on the expert policy can be replaced by practical (and often realistic) assumptions on available algorithms.

1.1 Contributions.

As in previous work, we formalize behavior cloning in two stages: at train-time , we learn a map from observations to distributions over actions, supervised by (state, action)-pairs from expert demonstrations coming from N exp independent expert trajectories, while at test-time , the learned map, or policy , is executed on random initial states (distributed identically to initial training states). Following the schematic of existing theoretical analyses of behavior cloning [68, 50, 27], we demonstrate that a policy trained by minimizing a certain supervised learning objective on expert demonstrations induces trajectories that approximate those of expert demonstrations. Our work considers a significantly more general setting than past theoretical literature, and one which reflects the strength of generative models for imitation. One corollary of our key contributions is summarized in the following informal statement. The main technical insights leading to the proof of the theorem are detailed in the bullet points below it, and depicted in Figure 1.

Theorem (informal). Consider a generative behavior cloner ˆ π that learns to predict sequences of expert actions on horizon H , along with low-level controllers that locally stabilize the trajectories. Then, with a suitable data noising strategy, for all times h ≤ H ,

where E expert ,h [ W 1 ( P expert actions , P imitator actions )] denotes a 1 -Wasserstein distance in an appropriate metric between the conditional distribution over expert and imitator actions given the observation at time step h , and where O ISS hides constants depending polynomially on the stability properties of the low-level controllers, defined formally in Section 3.1.

We now detail the key ingredients of our results.

We imitate stochastic demonstrators that may exhibit both complex correlations between actions in their trajectories (e.g. be non-Markovian) and multi-modal behavior. The natural object to imitate in this setting is the conditional probability distribution of expert actions given recent states, but marginalized over past states. We require said conditional action distribution to be learnable by a generative model , but otherwise arbitrarily complex: in particular, the conditional distribution of an expert actions given the state can be discontinuous (in any natural distance metric) as a function of state, as in the bifurcation depicted in Figure 1 (right) .

We obtain rigorous, theoretical guarantees and without requiring either interactive data collection (e.g. DAGGER [57, 40]), or access to gradients of the expert policy (as in TASIL[50]). Instead, we replace these assumption with an oracle, described below, which synthesizes stabilizing, low-level policies along training demonstrations-the green arrows in Figure 1 (left) . This mirrors recent work on generative behavior cloning that find that providing state-commands through inverse dynamics controllers [33, 2] or position-command controllers of end effectors [19] leads to substantially improved performance.

We also apply a subtle-yet-significant modification to a popular data noising strategy, which we show yields both theoretical and empirical beneifts. Data noising ensures a helpful property we denote total variation continuity that interpolates between modes in probability space (without naively averaging their trajectories in world space). This effectively 'fills in the missing gaps' in bifurcations, as indicated by yellow arrows in Figure 1.

Our main results, Theorems 1 and 2, are reductions from imitation of complex expert trajectories to supervised generative learning of a specific conditional distribution. For concreteness, Theorem 3 instantiates the generative modeling with Denoising Diffusion Probabilistic Models (DDPMs) of sufficient regularity and expressivity (as investigated empirically in [19, 48, 26]), and establishes end-to-end guarantees for imitation of complex trajectories with sample complexity polynomial in relevant problem parameters. Our analysis framework exposes that any sufficient powerful generative learner obtains similar guarantees. Finally, we empirically validate the benefits of our proposed smoothing strategy in simulated robotic manipulation tasks. We now summarize the algorithmic choices and analytic ideas that facilitate our reduction.

Abridged Related Work. Due to space, we defer a full comparison to past work to Appendix B. DDPMs, proposed in [29, 60], along with their relatives have seen success in image generation [62, 55], along with imitation learning (without data augmentation) [33, 19, 48], which is the starting point of our work. Smoothing data augmentation is ubiquitous in modern imitation learning [40] and our approach corresponds to that of [36] but with noise added at inference time. Despite the benefits of adaptive data collection [58, 40], adaptive demonstrations are more expensive to collect. Previous analyses of imitation learning without adaptive data collection have focused on classical controltheoretic notions of stability, notably incremental stability, [68, 27, 50], which require continuity, Markovianity, and often determinism, and preclude the bifurcations permitted in our setting.

2 Setting

Notation and Preliminaries. Appendix A gives a full review of notation. Bold lower-case (resp. upper-case) denote vectors (resp. matrices). We abbreviate the concatenation of sequences via z 1: n = ( z 1 , . . . , z n ) . Norms ∥ · ∥ are Euclidean for vectors and operator norms for matrices unless

otherwise noted. Rigorous probability-theoretic preliminaries are provided in Appendix F. In short, all random variables take values in Polish spaces X (which include real vector spaces), the space of Borel distributions on X is denoted ∆ ( X ) . We rely heavily on couplings from optimal transport theory: given measures X ∼ P and X ′ ∼ P ′ on X and X ′ respectively, C ( P , P ′ ) denotes the space of joint distributions µ ∈ ∆ ( X × X ′ ) called 'couplings' such that ( X,X ′ ) ∼ µ has marginals X ∼ P and X ′ ∼ P . ∆ ( X | Y ) denotes the space of conditional probability distributions Q : Y → ∆ ( X ) , formally called probability kernels ; Appendix F rigorously justifies that, in our setting, all conditional distributions can be expressed as kernels (which we do throughout the paper without comment). Finally I ∞ ( E ) denotes the indicator taking value 1 if E is true and ∞ otherwise.

Dynamics and Demonstrations. Weconsider a discrete-time, control system with states x t ∈ X := R d x , and inputs u t ∈ U := R d u , obeying the following nonlinear dynamics

Given length T ∈ N , we call sequences ρ T = ( x 1: T +1 , u 1: T ) ∈ P T := X T +1 ×U T trajectories . For simplicity, we assume that (2.1) is deterministic and address stochastic dynamics in Appendix N. Though the dynamics are Markov and deterministic, we consider a stochastic and possibly nonMarkovian demonstrator, which allows for the multi-modality described in the Section 1.

Definition 2.1 (Expert Distribution) . Let D exp ∈ ∆ ( P T ) denote an expert distribution over trajectories to be imitated. D x 1 denotes the distribution of x 1 under ρ T = ( x 1: T +1 , u 1: T ) ∼ D exp .

Primitive Controllers. Our approach is to imitate not just actions, but simple local control policies . In the body of this paper, we consider affine mappings X → U (redundantly) parameterized as x ↦→ ¯ u + ¯ K ( x -¯ x ) ; we call these primitive controllers , denoted with κ = (¯ u , ¯ x , ¯ K ) ∈ K . We describe the synthesis of these controllers in Appendix D , and extend our results to general families of parameterized controllers in Appendix E. We argue in Appendix E that primitive controllers are in fact standard practice, and implicit via robotic position control in many applications of diffusion to robotic behavior cloning.

Chunking Policies and Indices. The expert distribution D exp may involve non-Markovian sequences of actions. We imititate these sequences via chunking policies . Fix a chunk length τ chunk ∈ N and observation length τ obs ≤ τ chunk , and define time indices t h = ( h -1) τ chunk +1 . For simplicity, we assume τ chunk divides T , and set H = T/τ chunk . Given a ρ T ∈ P T , define the trajectory-chunks and observation chunks

for h > 1 , and s 1 = o 1 = x 1 (for simplicity, we embed o 1 into P τ obs -1 via zero-padding). We call τ chunk -length sequences of primitive controllers composite actions

A chunking policy π = ( π h ) consists of functions π h mapping observation-chunks o h to distributions ∆ ( A ) over composite actions and interacting with the dynamics (2.1) by a h = κ t h : t h -1 ∼ π h ( o h ) , and executing u t = κ t ( x t ) . We let d A = τ chunk ( d x + d u + d x d u ) denote the dimension of the space A of composite actions. The chunking scheme is represented in Figure 2, demonstrating the rationale for using primitive controllers over open-loop actions. Remark C.1 describes our rationale for studying states over generic observations, and considering time-dependent policies.

Desideratum. The quality of imitation of a deterministic policy is naturally measured in terms of step-wise closeness of state and action [68, 50]. With stochastic policies, however, two rollouts of even the same policy can visit different states. We propose measuring distributional closeness via couplings introduced in the preliminaries above. We define the following losses, focusing on the marginal distributions between trajectories.

Definition 2.2. Given ε > 0 and a (chunking) policy π , the (marginal distribution) imitation loss is

L marg ,ε ( π ) := max t ∈ [ T ] inf µ P µ [ max { ∥ x exp t +1 -x π t +1 ∥ , ∥ u exp t -u π t ∥ } > ε ] , where the infimum is over all couplings µ between the distribution of ρ T under D exp and that induced by the policy π as described above, such that P µ [ x exp 1 = x π 1 ] = 1 .

̂

x t

x t

-

1

u t

:

t

+

τ

chunk

-

1

∼

π

(

x t

-

τ

obs

:

t

)

x t

-

τ

obs

u t

x t

+1

u t

+1

x t

+2

x t

+

τ

chunk

x

⋆

t

+1

x

⋆

t

+2

x

⋆

t

+

τ

chunk

x t

x t

-

1

k t

:

t

+

τ

chunk

-

1

∼

π

(

x t

-

τ

obs

:

t

)

x t

-

τ

obs

k t

(

x t

)

x t

+1

x t

+2

x t

+

τ

chunk

x

⋆

t

+1

x

⋆

t

+2

x

⋆

t

+

τ

chunk

k t

+1

(

x t

+1

)

Expert Trajectory

Action chunking

Controller chunking

̂

̂

̂

̂

̂

Under stronger conditions (whose necessity we establish), we can also imitate joint distributions over actions (Appendix J). Observe that L fin ,ε ≤ L marg ,ε , and that both losses are equivalent to Wasserstein-type metrics on bounded domains. These losses are also equivalent to Lévy-Prokhorov metrics [64] under re-scaling of the Euclidean metric (even for unbounded domains), and also correspond to total variation analogues of shifted Renyi divergences [5, 6]. While empirically evaluating these infima over couplings is challenging, L marg ,ε upper bounds the difference in expectation between any bounded and Lipschitz control cost decomposing across time steps, states and inputs, and L fin ,ε upper bounds differences in final-state costs; see Appendix J for discussion.

Diffusion Models. Our analysis provides imitiation guarantees when chunking policies π h select a h via a sufficiently accurate generative model. Given their recent success, instantiate our analysis for the popular Denoising Diffusion Probabilistic Models (DDPM) framework [18, 41] that allows the learner to sample from a density q ∈ ∆( R d ) assuming that the score ∇ log q is known to the learner. More precisely, suppose the learner is given an observation o h and wishes to sample a h ∼ q ( ·| o h ) for some family of probability kernels q ( ·|· ) . ADDPMstarts with some a 0 h sampled from a standard Gaussian noise and iteratively 'denoises' for each DDPM-time step 0 ≤ j < J :

where s θ,h ( a j h , o h , j ) estimates the true score s ⋆,h ( a h , o h , αj ) , formally defined for any continuous argument t ≤ Jα to be s ⋆,h ( a , o h , t ) := ∇ a log q ⋆ h, [ t ] ( a | o h ) , where q ⋆ h, [ t ] ( ·| o h ) is the distribution of e -t a (0) h + √ 1 -e -2 t γ with a (0) h is sampled from the target distribution we which to sample from, and γ ∼ N (0 , I ) is a standard Gaussian. We denote by DDPM ( s θ , o h ) the law of a J h sampled according to the DDPM using s θ ( · , o h , · ) as a score estimator. Preliminaries on DPPMs are detailed in Appendix L.

3 Conditional sampling with stabilization suffices for behavior cloning

We show that trajectories of the form given in Definition 2.1 can be efficiently imitated if (a) we are given a synthesis oracle , described below, that produces low-level control policies that locally stabilize chunks of the trajectory with primitive controllers and (b) we can learn to generate certain appropriate distributions over composite actions, i.e. sequences of primitive controllers. All the following results apply to affine primitive controllers introduced in Section 2 and assume that the system dynamics are second-order smooth and locally stabilizable. In Appendix E, we show that our results still hold with general families of parametric primitive controllers, provided that these controllers induce the same local stability guarantee.

The synthesis oracle. We say primitive controller (cf. Section 2) κ 1: T ∈ K T is consistent with a trajectory ρ = ( x 1: T +1 , u 1: T ) ∈ P T if ¯ x t = x t and ¯ u t = u t for all t ∈ [ T ] ; note that this implies that κ t ( x t ) = u t for all t . A synthesis oracle synth maps P T → K T such that, for all ρ T ∈ P T , κ 1: T = synth ( ρ T ) is consistent with ρ T . For our theory, we assume access to a synthesis oracle at training time, and assume the ability to estimate conditional distributions over joint sequences of primitive controllers; Appendix K explains how this can be implemented by solving Ricatti equations if dynamics are known (e.g. in a simulator), smooth, and stabilizable. In our experimental environment, control inputs are desired robot configurations, which the simulated

̂

̂

̂

robot executes by applying feedback gains. As discussed in Appendix E, learned or hand-coded low-level controllers are popular in practical implementations of generative behavior cloning. We discuss the merits of studying imitation learning with a synthesis oracle in depth in Appendix C.3.

Notions of distance. While restricting ourselves to affine primitive controllers, our approximation error of generative behavior cloner is measured in terms of optimal transport distances that use the following 'maximum distance.' Given two composite actions a = (¯ u 1: τ chunk , ¯ x 1: τ chunk , ¯ K 1: τ chunk ) and a ′ = (¯ u ′ 1: τ chunk , ¯ x ′ 1: τ chunk , ¯ K ′ 1: τ chunk ) , we define

Distances between policies are defined via natural optimal transport costs. Given two policies π = ( π h ) , π ′ = ( π ′ h ) and observation chunk o h , we define an induced optimal transport cost

where the inf µ denotes the infinum over all couplings between a h ∼ π h ( o h ) and a ′ h ∼ π ′ h ( o h ) . ∆ ε corresponds to a relaxed Lévy-Prokhorov metric [64], and can always be bounded, via Markov's inequality, by ∆ ε ( π h , π ′ h | o h ) ≤ 1 ε W 1 , d max ( π h ( o h ) , π ′ h ( o h )) , where W 1 , d max ( π h ( o h ) , π ′ h ( o h )) denotes the 1 -Wasserstein distance between a h ∼ π h ( o h ) and a ′ h ∼ π ′ h ( o h ) .

3.1 Incremental Stability and the Synthesis Oracle.

We assume that synthesis oracle above produces incrementally stabilizing control gains, in the sense first proposed by [9]. Incremental stability has emerged as a natural desirable property for imitation limitation [50, 68, 27], because it forces the expert to be robust to small perturbations of their policy. We now supply a formal definition. Given a primitive controller κ : R d x → R d u , define the closed loop dynamic map f cl ,κ ( x , δ u ) := f ( x , κ ( x ) + δ u ) . Thus, composite action a is consistent with a trajectory chunk s = ( x 1: τ chunk +1 , u 1: τ ) if x t +1 = f cl ,κ t ( x , 0 ) for 1 ≤ t ≤ τ chunk . 1

Definition 3.1 (Time-Varying Incremental Stability) . Let γ ( · ) be a class K function, β ( · , · ) be class KL function, and let a = ( κ 1 , κ 2 , . . . , κ τ ) denote a sequence of primitive controllers (i.e. a composite action when τ = τ chunk ). Given a sequence of input perturbations δ u 1: τ ∈ ( R d u ) τ and initial condition ξ ∈ R d x , let x a i +1 ( δ u 1: τ , ξ ) = f cl ,κ i ( x a i ( δ u 1: τ , ξ ) , δ u i ) , with x a 1 = ξ . We say that composite action a is time-varying incrementally input-to-state stable ( t -ISS) with moduli γ ( · ) , β ( · , · ) if for all ξ , ξ ′ ∈ R d x , 0 ≤ i ≤ τ , ∥ x a i ( 0 1: τ , ξ ) -x a i ( δ u 1: τ , ξ ′ ) ∥ ≤ β ( ∥ ξ -ξ ′ ∥ , τ ) + γ (max 1 ≤ s ≤ i -1 ∥ δ u s ∥ ) . Given parameters c γ , c ξ > 0 we say that a is localt -ISS at ξ 0 if the above holds only for all ξ , ξ ′ , δ u 1: τ such that ∥ ξ -ξ 0 ∥ , ∥ ξ ′ -ξ 0 ∥ ≤ c ξ and max t ∥ δ u t ∥ ≤ c γ .

Incremental stability implies that as the inital conditions ∥ ξ -ξ ′ ∥ → 0 and max 0 ≤ s ≤ i -1 ∥ δ u t ∥ → 0 , the trajectories induced by taking rolling out a from ξ , and rolling out a from ξ ′ with additive input perturbations δ u 1: τ tend to zero in norm. This behavior needs only hold for initial conditions in a small neighborhood of a nominal state ξ 0 . Importantly, the perturbations δ u 1: τ are fixed pertubrations of inputs, applied to the closed loop behavior under the controllers. Our notion of incremental stability are similar too, but sublty different similar notions of past work. We provide an extended comparisons in Appendix E.2. Our main assumption is that the synthesis oracle described above produced primitive controllers which are consistent with, and incrementally stabilizing for, the demonstrated trajectories. Figure 1 demonstrates the effect of stabilizing primitive controllers.

Assumption 3.1. We assume that our synthesis oracle enjoys the following property. Let ρ T = ( x 1: T +1 , u 1: T ) ∼ D exp , and let κ 1: T = synth ( ρ T ) , partitioned into composite actions a 1: H , with κ t ( x ) = ¯ K t ( x -¯ x t ) + ¯ u t . We assume that, with probability one, κ 1: T is consistent with ρ T 2 , and that, for each 1 ≤ h ≤ H , a h = ( κ t h : t h + τ chunk -1 ) is local t -ISS at x t h with moduli γ , β and parameters c β , c ξ > 0 . We further assume that γ and β take the form

(a)

X

X

(b)

Mostly Left

Equal Split

Mostly Right

Right Mode

Lastly, we assume that for the expert trajectories and the primitive controllers drawn as above, it holds that satisfy max {∥ x t ∥ , ∥ u t ∥} ≤ R dyn and ∥ ¯ K t ∥ ≤ R K with probability one.

In Appendix K, we show that Assumption 3.1 holds whenever (a) the dynamics of our system are smooth (but not necessarily linear!) (b) the affine gains are chosen to stabilize the Jacobian linearizations of the system around the nominal trajectory.

Definition 3.2 (Problem constants) . Throughout, we refer to constants c 1 , c 2 , c 3 , c 4 , c 5 > 0 , which are polynomial in the terms in Assumption 3.1, and which are defined formally in Appendix K.

3.2 Simplified guarantees under total variation continuity

This section presents our main theoretical result: if one learns a chunking policy ˆ π that can compute the conditional distribution of composite actions at time steps given observation-chunks, then a stochastically smoothed version of this policy, ˆ π σ , has low imitation error. Define, for any length τ ∈ N , the trajectory distance between trajectories ρ = ( x 1: τ +1 , u 1: τ ) , ρ ′ = ( x ′ 1: τ +1 , u ′ 1: τ ) ∈ P τ

In particular, we define d traj ( o h , o ′ h ) and d traj ( s h , s ′ h ) by viewing these as trajectories of length τ obs -1 and τ chunk , respectively. Lastly, we define a per-timestep restriction of the expert distribution. In this section, we consider the case where the learner policy satisfies a total variation continuity (TVC) condition, defined below.

Definition 3.3 (TVC of Chunking Policies) . We say that a chunking policy π = ( π h ) is total variation continuous with modulus γ TVC : R ≥ 0 → R ≥ 0 , written γ TVC -TVC, if, for all h ∈ [ H ] and any observation-chunks o h , o ′ h ∈ P τ obs -1 , TV ( π h ( o h ) , π h ( o ′ h )) ≤ γ TVC ( d traj ( o h , o ′ h )) .

Wedepict the TVC property using our running left-right obstacle example in Figure 3. We stress that, in Definition 3.3, the TV bound on TV ( π h ( o h ) , π h ( o ′ h )) applies to the composite actions consisting of primitive controllers a h = κ t h : t h + τ chunk -1 ∼ π h ( o h ) ; it does not upper bound the TV distance between raw control inputs. Indeed, ensuring TVC of the latter can lead to the failure modes depicted in Figure 3(b). Next, we extract an expert 'policy' from the expert demonstrations.

Definition 3.4 (Expert 'policy' with synthesized controllers) . For h ∈ [ H ] , we let D exp ,h denote the joint distribution of ( a h , o h ) , induced by drawing a trajectory ρ T = ( x 1: T +1 , u 1: T ) ∼ D exp from the expert distribution, κ 1: T = synth ( ρ T ) be the associated primitive controllers, letting o h = ( x t h -τ obs +1: t h , u t h -τ obs +1: t h -1 ) be the associated observation-chunk at time h , and a h = κ t h : t h +1 -1 the associated composite action. We let π ⋆ h ( · ) : O → ∆ ( A ) denote the condition distribution of a h | o h under D exp ,h .

The conditional distributions π ⋆ h ( · ) are estimated when training a generative model to predict a h from observations o h . Note that π ⋆ h ( · ) (and D exp ,h ) is defined in terms of both expert demonstration

from D exp and the associated synthesized primitive controllers. In Lemma J.6, we show that when the synthesis oracle κ 1: T = synth ( ρ T ) produces primitive controllers consistent with the trajectories, than π ⋆ = ( π ⋆ h ) produces the same marginals over states as D exp ; that is, L marg ,ε ( π ⋆ ) = 0 .

Theorem 1. Suppose Assumption 3.1 holds, and suppose that 0 ≤ ε < c 2 , and τ chunk ≥ c 3 . Then, for any non-decreasing non-negative γ TVC ( · ) and γ TVC-TVC chunking policy ˆ π , it holds that L marg ,ε (ˆ π ) ≤ Hγ TVC ( ε ) + ∑ H h =1 E o h ∼D exp ,h ∆ ( ε/c 1 ) ( π ⋆ h ( o h ) , ˆ π h ( o h )) , which is at most Hγ TVC ( ε ) + c 1 ε ∑ H h =1 E o h ∼D exp ,h [ W 1 , d max ( π ⋆ h ( o h ) , ˆ π h ( o h ))] .

The above result reduces the marginal imitation error of ˆ π to the sum over optimal transport errors between ˆ π and a | o h chosen by the expert demonstrators. Thus, if these are small, the local stabilization properties of the primitive controllers guaranteed by Assumption 3.1 ensure that errors compound at most linearly in problem horizon. The key ideas of the proof are given Appendix D, via a general template for imitation learning of general stochastic policies. This template is instantiated with a details in Appendix J.

3.3 A general guarantee via data noising.

To circumvent assuming that the learner's policy is TVC, we study estimating the conditionals under a popular data augmentation technique [36], where the learner is trained to imitate the conditional sequence of a | ˜ o h , where ˜ o h ∼ N ( o h , σ 2 I ) adds σ 2 -variance Gaussian noise to the true observation-chunk. To understand this better, consider the following smoothed policy:

Definition 3.5 (The smoothed policy) . Let ˆ π = (ˆ π h ) be a chunking policy. We define the smoothed policy ˆ π σ = (ˆ π σ,h ) by letting ˆ π σ,h ( · | o h ) be distributed as ˆ π h ( · | ˜ o h ) , where ˜ o h ∼ N ( o h , σ 2 I ) .

Appendix J.7.2 show's that Pinsker's inequality implies noising automatically enforces TVC This suggests that we can use some form of data noising to enforce the TVC property in Definition 3.3. Let's now consider a related problem: trying to estimate the optimal distribution over composite actions conditioned on a noised observation. This gives rise to a deconvolution of the expert policy, which can be thought as an inverse operation of data noising.

Definition 3.6 (Noised Data Distribution and Deconvolution Policy) . Let D exp ,h be as in Definition 3.4. Define D exp ,σ,h as the distribution over (˜ o h , a h ) generated by ( o h , a h ) ∼ D exp ,h and ˜ o h ∼ N ( o h , σ 2 I ) . We define the deconvolution policy π ⋆ dec ,σ,h (˜ o h ) as the conditional distribution of a h | ˜ o h under D exp ,σ,h .

Analogously to π ⋆ , the policy π ⋆ dec ,σ,h is what a generative model trained to generate a h from noised observations ˜ o h of o h ∼ D exp learns to generate. Our next theorem states that, if our ˆ π approximates the idealized conditional distributation of composite actions given noised observations, then ˆ π σ , the smoothed policy, imitates the expert distribution with provable bounds on its imitation error:

Theorem 2 (Reduction to conditional sampling under nosing) . Suppose Assumption 3.1 holds. Let c 1 , . . . , c 5 > 0 , defined in Definition 3.2, and let Θ Iss ( x ) denote a term which is upper and lower bounded by a x times a polynomial in those constants and their inverses. Then, for ε ≤ Θ Iss (1) , if we choose σ = ε/ Θ Iss ( √ d x +log(1 /ε )) and let τ chunk ≤ c 3 and τ chunk -τ obs ≥ 1 L β log( c 1 /ε ) ,

To reiterate, Theorem 2 guarantees imitation of the distribution of marginals and final states of D exp by replacing the explicit TVC assumption with noising, and the resulting guarantee applies to the smoothed policy ˆ π σ which adds smoothing noise back in. Appendix J gives a number of additional results. In Appendix I, we show that the proof framework, outlined in Appendix D, which under lies the proofs of Theorems 1 and 2, is essentially sharp in the worst case. Moreover, in Appendix C.3, we discuss the merits and drawbacks of our use of the synthesis oracle, and how it circumvents some of the challenges encountered in behavior cloning in past work. The key intuition behind the proof of Theorem 2 is depicted in Figure 4, and full proof sketch is deferred to Appendix C.2

(a)

(b)

Right Mode

Trajectory

Noised State

4 HINT: Instantiating Data Noising with DDPMs

We now instantiate Theorem 2 by showing that one can learn a policy ˆ π for which the error terms in Theorems 1 and 2 are small by fitting a DDPM to noise-smoothed data. Our proposed algo-

Algorithm 1 H ierarchical I mitation via N oising at Inference T ime (HINT)

[h]

rithm, HINT (Algorithm 1) combines DDPM-learning of chunked policies as in [19] with a popular form of data-augmentation [36]. We collect N exp expert trajectories, synthesize gains, and segment trajectories into observation-chunks o h and composite actions a h as described in Section 2. We perturb each o h to form N aug chunks ˜ o h , as well as horizon indices j ∈ [ J ] and inference noises γ ∼ N (0 , ( αj h ) 2 I ) , and add these tuples ( a h , ˜ o h , j h , γ h , h ) to our data D . We end the training phase by minimizing the standard DDPM loss [62]:

Our algorithm differs subtly from past work in Line 8: motivated by Theorem 2, we add smoothing noise back in at test time. Here, the notation DDPM ( s θ,h , · ) ◦N ( o h , σ 2 I ) means, given o h , we perturb it to ˜ o h ∼ N ( o h , σ 2 I ) , and sample a h ∼ DDPM ( s θ,h , ˜ o h ) . We now state an informal guarantee for HINT, deferring a formal statement to Appendix C.5.

Theorem (Informal Theorem) . Suppose that the system dynamics are smooth and that Assumption 3.1 holds for the linearized system. Then there is a choice of the parameters in HINT that is polynomial in all problem parameters such that for N exp , polynomially large in problem parameters, L marg ,ε ( ̂ π σ ) ≤ Θ ( εH √ τ obs ( √ d x +log(1 /ε )) ) with high probability.

4.1 Experimental Results

In this section, we demonstrate the benefits of diffusing low level controllers, and of our approach to data noising. We explain the environments in greater detail, along with all training and compu-

0.00

0.05

0.10

0.15

0.20

0.25

0.30

0.0

0.2

0.4

0.6

0.8

1.0

Normalized Reward

No Gains, N=1

No Gains, N=5

Gains, N=1

Gains, N=3

Gains, N=5

Gains, N=10

1

2

3

4

5

6

7

8

9

10

Number of Trajectories

0.0

0.2

0.4

0.6

0.8

1.0

Normalized Reward

No Gains,

=0.0

No Gains,

=0.02

No Gains,

=0.1

No Gains,

=0.3

Gains,

=0.0

Gains,

=0.02

Gains,

=0.1

Gains,

=0.3

-̸-T-X-D-U-H-∗-↔-V-V-H-P-E-O-∩

-⊇-≡-⊇-∼

-⊇

-⊇-≡-≤

-⊇-≡-≥

-⊇-≡-⪯

-⊇-≡-⪰

-⊇

-⊇-≡-≥

-⊇

-⊇-≡-⊇-≤

-⊇-≡-⊇-≥

-⊇-≡-⊇-⪯

-⊇-≡-⊇-⪰

-⊇-≡-⊇-∼

-⊇

-⊇-≡-⊇-≤

-⊇-≡-⊇-≥

-⊇-≡-⊇-⪯

-⊇-≡-⊇-⪰

-∞-R-L-V-H-∗-̸-L-J-P-D

-∞-R-L-V-H-∗-̸-L-J-P-D

-∞-R-L-V-H-∗-̸-L-J-P-D

-↘-D-Q-∗-∋-L-F-N-∗-D-Q-G-∗-∋-O-D-F-H

tational details in Appendix O. 3 Figure 5 compares the performance of diffusing (chunks of) raw control inputs to diffusing (chunks of) gain matrices for a canonical model of a 2-d quadrotor. We find that diffusing gain matrices yields dramatic improvements in performance, in particular allowing a single imitated trajectory to outperform learning raw control inputs from 10 demonstrations.

Next, empirically evaluate the effect on policy performance of our proposal to inject noise back into the dynamics at inference time. We consider three challenging robotic manipulation tasks studied in prior work: PushT block-pushing [19]; Robomimic Can Pick-and-Place and Square Nut Assembly [42] (we direct the reader to Chi et al. [19] for an extensive empirical investigation into the performance of diffusion policies in the un-noised σ = 0 regime). We display the results of our experiments in Figure 6. Observe that the performance degredation of the replica policy from the unnoised σ = 0 variant is minimal across all environments and even leads to a slight but noticeable improvement in the small-noise regime for PushT (and low-data Can Pick and Place). In the presence of non-negligible noise HINT significantly outperforms the conventional policy ˆ π (obtained by noising observations at training but not test time), as predicted by our theory.

5 Discussion

This work considerably loosened assumptions placed on the expert distribution by introducing a synthesis oracle responsible for stabilization. How best to achieve low-level stabilization remains an open question. We hope that this work encourages further empirical research into improving the stability of imitation learning, either via the hierarchical route proposed in this paper or via new innovations.

Acknowledgments

We would like to thank Cheng Chi for his extensive help in running the code in [19] and Benjamin Burchfiel for numerous insightful discussions and intuitions about DDPM policies. Lastly, we thank Alexandre Amice for his helpful feedback on an earlier draft of this manuscript. AB acknowledges support from the National Science Foundation Graduate Research Fellowship under Grant No. 1122374 as well as the support of DOE through award DE-SC0022199. MS and RT acknowledge support from Amazon.com Services LLC grant; PO# 2D-06310236. DP and AJ acknowledge support from the Office of Naval Research under ONR grant N00014-23-1-2299.

References

S. Abu-El-Haija, N. Kothari, J. Lee, P. Natsev, G. Toderici, B. Varadarajan, and S. Vijayanarasimhan. Youtube-8m: A large-scale video classification benchmark. arXiv preprint arXiv:1609.08675 , 2016.

A. Ajay, Y. Du, A. Gupta, J. Tenenbaum, T. Jaakkola, and P. Agrawal. Is conditional generative modeling all you need for decision-making? arXiv preprint arXiv:2211.15657 , 2022.

A. Ajay, A. Gupta, D. Ghosh, S. Levine, and P. Agrawal. Distributionally adaptive meta reinforcement learning. arXiv preprint arXiv:2210.03104 , 2022.

A. Ajay, S. Han, Y. Du, S. Li, A. Gupta, T. Jaakkola, J. Tenenbaum, L. Kaelbling, A. Srivastava, and P. Agrawal. Compositional foundation models for hierarchical planning. arXiv preprint arXiv:2309.08587 , 2023.

J. Altschuler and K. Talwar. Privacy of noisy stochastic gradient descent: More iterations without more privacy loss. Advances in Neural Information Processing Systems , 35:37883800, 2022.

J. M. Altschuler and S. Chewi. Faster high-accuracy log-concave sampling via algorithmic warm starts. arXiv preprint arXiv:2302.10249 , 2023.

B. D. Anderson and J. B. Moore. Optimal control: linear quadratic methods . Courier Corporation, 2007.

O. Angel and Y. Spinka. Pairwise optimal coupling of multiple random variables. arXiv preprint arXiv:1903.00632 , 2019.

D. Angeli. A lyapunov approach to incremental stability properties. IEEE Transactions on Automatic Control , 47(3):410-421, 2002.

M. Bansal, A. Krizhevsky, and A. Ogale. Chauffeurnet: Learning to drive by imitating the best and synthesizing the worst. arXiv preprint arXiv:1812.03079 , 2018.

P. L. Bartlett and S. Mendelson. Rademacher and gaussian complexities: Risk bounds and structural results. Journal of Machine Learning Research , 3(Nov):463-482, 2002.

P. L. Bartlett, A. Montanari, and A. Rakhlin. Deep learning: a statistical viewpoint. Acta numerica , 30:87-201, 2021.

A. Block, Y. Mroueh, and A. Rakhlin. Generative modeling with denoising auto-encoders and langevin sampling. arXiv preprint arXiv:2002.00107 , 2020.

A. Block, Y. Mroueh, A. Rakhlin, and J. Ross. Fast mixing of multi-scale langevin dynamics under the manifold hypothesis. arXiv preprint arXiv:2006.11166 , 2020.

M. Bojarski, D. Del Testa, D. Dworakowski, B. Firner, B. Flepp, P. Goyal, L. D. Jackel, M. Monfort, U. Muller, J. Zhang, et al. End to end learning for self-driving cars. arXiv preprint arXiv:1604.07316 , 2016.

J. D. Chang, M. Uehara, D. Sreenivas, R. Kidambi, and W. Sun. Mitigating covariate shift in imitation learning via offline data without great coverage. arXiv preprint arXiv:2106.03207 , 2021.

L. Chen, K. Lu, A. Rajeswaran, K. Lee, A. Grover, M. Laskin, P. Abbeel, A. Srinivas, and I. Mordatch. Decision transformer: Reinforcement learning via sequence modeling. Advances in neural information processing systems , 34:15084-15097, 2021.

S. Chen, S. Chewi, J. Li, Y. Li, A. Salim, and A. Zhang. Sampling is as easy as learning the score: theory for diffusion models with minimal data assumptions. In NeurIPS 2022 Workshop on Score-Based Methods , 2022.

C. Chi, S. Feng, Y. Du, Z. Xu, E. Cousineau, B. Burchfiel, and S. Song. Diffusion policy: Visuomotor policy learning via action diffusion. arXiv preprint arXiv:2303.04137 , 2023.

S. Dasari, F. Ebert, S. Tian, S. Nair, B. Bucher, K. Schmeckpeper, S. Singh, S. Levine, and C. Finn. Robonet: Large-scale multi-robot learning. arXiv preprint arXiv:1910.11215 , 2019.

P. De Haan, D. Jayaraman, and S. Levine. Causal confusion in imitation learning. Advances in Neural Information Processing Systems , 32, 2019.

R. Durrett. Probability: theory and examples , volume 49. Cambridge university press, 2019.

C. Finn, T. Yu, T. Zhang, P. Abbeel, and S. Levine. One-shot visual imitation learning via meta-learning. In Conference on robot learning , pages 357-368. PMLR, 2017.

D. J. Foster and A. Rakhlin. L-infinity vector contraction for rademacher complexity. arXiv preprint arXiv:1911.06468 , 6, 2019.

J. W. Hagood and B. S. Thomson. Recovering a function from a dini derivative. The American Mathematical Monthly , 113(1):34-46, 2006.

P. Hansen-Estruch, I. Kostrikov, M. Janner, J. G. Kuba, and S. Levine. Idql: Implicit q-learning as an actor-critic method with diffusion policies. arXiv preprint arXiv:2304.10573 , 2023.

A. Havens and B. Hu. On imitation learning of linear control policies: Enforcing stability and robustness constraints via lmi conditions. In 2021 American Control Conference (ACC) , pages 882-887. IEEE, 2021.

D. Hendrycks, S. Basart, N. Mu, S. Kadavath, F. Wang, E. Dorundo, R. Desai, T. Zhu, S. Parajuli, M. Guo, et al. Jacob steinhardt et justin gilmer. the many faces of robustness: A critical analysis of out-of-distribution generalization. arXiv preprint arXiv:2006.16241 , 2020.

J. Ho, A. Jain, and P. Abbeel. Denoising diffusion probabilistic models. Advances in Neural Information Processing Systems , 33:6840-6851, 2020.

A. Hussein, M. M. Gaber, E. Elyan, and C. Jayne. Imitation learning: A survey of learning methods. ACM Computing Surveys (CSUR) , 50(2):1-35, 2017.

A. Hussein, E. Elyan, M. M. Gaber, and C. Jayne. Deep imitation learning for 3d navigation tasks. Neural computing and applications , 29:389-404, 2018.

A. Hyvärinen and P. Dayan. Estimation of non-normalized statistical models by score matching. Journal of Machine Learning Research , 6(4), 2005.

M. Janner, Y. Du, J. B. Tenenbaum, and S. Levine. Planning with diffusion for flexible behavior synthesis. arXiv preprint arXiv:2205.09991 , 2022.

C. Jin, P. Netrapalli, R. Ge, S. M. Kakade, and M. I. Jordan. A short note on concentration inequalities for random vectors with subgaussian norm. arXiv preprint arXiv:1902.03736 , 2019.

L. P. Kaelbling and T. Lozano-Pérez. Hierarchical task and motion planning in the now. In 2011 IEEE International Conference on Robotics and Automation , pages 1470-1477. IEEE, 2011.

L. Ke, J. Wang, T. Bhattacharjee, B. Boots, and S. Srinivasa. Grasping with chopsticks: Combating covariate shift in model-free imitation learning for fine manipulation. In 2021 IEEE International Conference on Robotics and Automation (ICRA) , pages 6185-6191. IEEE, 2021.

M. Kelly, C. Sidrane, K. Driggs-Campbell, and M. J. Kochenderfer. Hg-dagger: Interactive imitation learning with human experts. In 2019 International Conference on Robotics and Automation (ICRA) , pages 8077-8083. IEEE, 2019.

H. Khalil. Nonlinear Systems . Pearson Education. Prentice Hall, 2002. ISBN 9780130673893. URL https://books.google.com/books?id=t_d1QgAACAAJ .

I. Kostrikov, D. Yarats, and R. Fergus. Image augmentation is all you need: Regularizing deep reinforcement learning from pixels. arXiv preprint arXiv:2004.13649 , 2020.

M. Laskey, J. Lee, R. Fox, A. Dragan, and K. Goldberg. Dart: Noise injection for robust imitation learning. In Conference on robot learning , pages 143-156. PMLR, 2017.

H. Lee, J. Lu, and Y. Tan. Convergence of score-based generative modeling for general data distributions. In International Conference on Algorithmic Learning Theory , pages 946-985. PMLR, 2023.

A. Mandlekar, D. Xu, J. Wong, S. Nasiriany, C. Wang, R. Kulkarni, L. Fei-Fei, S. Savarese, Y. Zhu, and R. Martín-Martín. What matters in learning from offline human demonstrations for robot manipulation. arXiv preprint arXiv:2108.03298 , 2021.

T. Marcucci, J. Umenberger, P. A. Parrilo, and R. Tedrake. Shortest paths in graphs of convex sets. arXiv preprint arXiv:2101.11565 , 2021.

A. Maurer. A vector-contraction inequality for rademacher complexities. In Algorithmic Learning Theory: 27th International Conference, ALT 2016, Bari, Italy, October 19-21, 2016, Proceedings 27 , pages 3-17. Springer, 2016.

M. Mezard and A. Montanari. Information, physics, and computation . Oxford University Press, 2009.

D. Misra. Mish: A self regularized non-monotonic activation function. arXiv preprint arXiv:1908.08681 , 2019.

A. Q. Nichol and P. Dhariwal. Improved denoising diffusion probabilistic models. In International Conference on Machine Learning , pages 8162-8171. PMLR, 2021.

T. Pearce, T. Rashid, A. Kanervisto, D. Bignell, M. Sun, R. Georgescu, S. V. Macua, S. Z. Tan, I. Momennejad, K. Hofmann, et al. Imitating human behaviour with diffusion models. arXiv preprint arXiv:2301.10677 , 2023.

E. Perez, F. Strub, H. De Vries, V. Dumoulin, and A. Courville. Film: Visual reasoning with a general conditioning layer. In Proceedings of the AAAI Conference on Artificial Intelligence , volume 32, 2018.

D. Pfrommer, T. Zhang, S. Tu, and N. Matni. Tasil: Taylor series imitation learning. Advances in Neural Information Processing Systems , 35:20162-20174, 2022.

D. Pfrommer, M. Simchowitz, T. Westenbroek, N. Matni, and S. Tu. The power of learned locally linear models for nonlinear policy optimization. arXiv preprint arXiv:2305.09619 , 2023.

Y. Polyanskiy and Y. Wu. Information Theory: From Coding to Learning . Cambridge University Press, 2022+.

M. Raginsky, A. Rakhlin, and M. Telgarsky. Non-convex learning via stochastic gradient langevin dynamics: a nonasymptotic analysis. In Conference on Learning Theory , pages 16741703. PMLR, 2017.

A. Rakhlin, K. Sridharan, and A. B. Tsybakov. Empirical entropy, minimax regret and minimax risk. Bernoulli Society for Mathematical Statistics and Probability , 2017.

A. Ramesh, P. Dhariwal, A. Nichol, C. Chu, and M. Chen. Hierarchical text-conditional image generation with clip latents. arXiv preprint arXiv:2204.06125 , 2022.

O. Ronneberger, P. Fischer, and T. Brox. U-net: Convolutional networks for biomedical image segmentation. In Medical Image Computing and Computer-Assisted Intervention-MICCAI 2015: 18th International Conference, Munich, Germany, October 5-9, 2015, Proceedings, Part III 18 , pages 234-241. Springer, 2015.

S. Ross and D. Bagnell. Efficient reductions for imitation learning. In Proceedings of the thirteenth international conference on artificial intelligence and statistics , pages 661-668. JMLR Workshop and Conference Proceedings, 2010.

S. Ross, G. Gordon, and D. Bagnell. A reduction of imitation learning and structured prediction to no-regret online learning. In Proceedings of the fourteenth international conference on artificial intelligence and statistics , pages 627-635. JMLR Workshop and Conference Proceedings, 2011.

N. M. Shafiullah, Z. Cui, A. A. Altanzaya, and L. Pinto. Behavior transformers: Cloning k modes with one stone. Advances in neural information processing systems , 35:22955-22968, 2022.

J. Sohl-Dickstein, E. Weiss, N. Maheswaranathan, and S. Ganguli. Deep unsupervised learning using nonequilibrium thermodynamics. In International Conference on Machine Learning , pages 2256-2265. PMLR, 2015.

J. Song, C. Meng, and S. Ermon. Denoising diffusion implicit models. arXiv preprint arXiv:2010.02502 , 2020.

Y. Song and S. Ermon. Generative modeling by estimating gradients of the data distribution. Advances in neural information processing systems , 32, 2019.

Y. Song, S. Garg, J. Shi, and S. Ermon. Sliced score matching: A scalable approach to density and score estimation. In Uncertainty in Artificial Intelligence , pages 574-584. PMLR, 2020.

V. Strassen. The existence of probability measures with given marginals. The Annals of Mathematical Statistics , 36(2):423-439, 1965.

X. Sun, S. Yang, and R. Mangharam. Mega-dagger: Imitation learning with multiple imperfect experts. arXiv preprint arXiv:2303.00638 , 2023.

R. Tedrake. Lqr-trees: Feedback motion planning on sparse randomized trees. 2009.

R. Thoppilan, D. De Freitas, J. Hall, N. Shazeer, A. Kulshreshtha, H.-T. Cheng, A. Jin, T. Bos, L. Baker, Y. Du, et al. Lamda: Language models for dialog applications. arXiv preprint arXiv:2201.08239 , 2022.

S. Tu, A. Robey, T. Zhang, and N. Matni. On the sample complexity of stability constrained imitation learning. In Learning for Dynamics and Control Conference , pages 180-191. PMLR, 2022.

R. Van Handel. Probability in high dimension. Technical report, PRINCETON UNIV NJ, 2014.

A. Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A. N. Gomez, Ł. Kaiser, and I. Polosukhin. Attention is all you need. Advances in neural information processing systems , 30, 2017.

R. Vershynin. High-dimensional probability: An introduction with applications in data science , volume 47. Cambridge university press, 2018.

C. Villani. Topics in optimal transportation , volume 58. American Mathematical Soc., 2021.

C. Villani et al. Optimal transport: old and new , volume 338. Springer, 2009.

P. Vincent. A connection between score matching and denoising autoencoders. Neural computation , 23(7):1661-1674, 2011.

M. J. Wainwright. High-dimensional statistics: A non-asymptotic viewpoint , volume 48. Cambridge university press, 2019.

T. Westenbroek, M. Simchowitz, M. I. Jordan, and S. S. Sastry. On the stability of nonlinear receding horizon control: a geometric perspective. In 2021 60th IEEE Conference on Decision and Control (CDC) , pages 742-749. IEEE, 2021.

Y. Wu and K. He. Group normalization. In Proceedings of the European conference on computer vision (ECCV) , pages 3-19, 2018.

T. Zhang, Z. McCarthy, O. Jow, D. Lee, X. Chen, K. Goldberg, and P. Abbeel. Deep imitation learning for complex manipulation tasks from virtual reality teleoperation. In 2018 IEEE International Conference on Robotics and Automation (ICRA) , pages 5628-5635. IEEE, 2018.

S. Zhao and A. Grover. Decision stacks: Flexible reinforcement learning via modular generative models. arXiv preprint arXiv:2306.06253 , 2023.

T. Z. Zhao, V. Kumar, S. Levine, and C. Finn. Learning fine-grained bimanual manipulation with low-cost hardware. arXiv preprint arXiv:2304.13705 , 2023.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper presents a simple hierarchical approach for non-Markovian imitation learning along with a rigorous theoretical analysis.\nThe method trains a high level policy, given by a Denoising Diffusion Probabilistic Model (DDPM), to predict a sequence of linear control gains based on a state-action history.\nThe DDPM is trained in supervised fashion (using the standard DDPM loss) where the labels (the controller gains) are computed from the given expert trajectories using an oracle (e.g. based on linearizations along the trajectory). To address covariate shift (compounding errors) in imitation learning the method applies a form of data augmentation by adding noise to the state-action history. In contrast to prior work, these perturbations are also applied during inference, which is theoretically motivated and empirically shown to improve imitation performance.\nThe main contributions seems to be the thorough analysis of this method, which leads to Theorem 1, a bound on the imitation performance. \nThe applicability of the proposed method is demonstrated on simulated robot tasks, although no comparisons with prior work are provided.\n\nSTRENGTHS:\nNovelty & Relevance\n-------------------\nBoth the method and in particular the analysis seem to be novel and non-Markovian imitation learning is an important field of research. I would not expect the proposed method to outperform other (theoretically less well understood) methods, but I do think that the general approach of using a diffusion model as a high-level policy to predict controller gains could be used by future work that is more focused on practical performance. Furthermore, the theoretical analysis could be valuable for theorists in the field of imitation learning. However, I am not very familiar with related theoretical work and cannot assess well, which aspects/techniques of the analysis might be most relevant.\n\nSoundness\n---------\nThe algorithm seems sounds. I guess one could argue that it might be simpler and more stable if the higher-level policy would predict a trajectory segment instead of the gains, given that we assume access to an oracle that can provide us with stabilizing gains. However, querying the oracle could be too costly, so directly predicting the gains can be sensible.\n\nI cannot fully confirm the accuracy of the theoretical results because verifying the complete derivations would take significant effort. However, the claims (including Theorem 1) seem reasonable, apart from non-critical issues (see Questions below).\n\nClarity\n-------\nThe paper is well-written. The paper is very strict in terms of notation, definitions, and statements, which increases clarity by reducing ambiguities.\n\nWEAKNESSES:\nClarity\n-------\nAlthough I mentioned the rigor as a strength, it also feels a bit pedantic. Furthermore, it is cumbersome to keep track of the different constants and definitions making the paper hard to read and making it easy to get lost in details.\n\nEvaluation\n----------\nDespite the fact that the contributions are mainly in the theoretical analysis, quantitative comparisons with alternate methods that have been used in these environments would be useful. I don't know which prior work was used in that environment that also provided code, but I assume it should be possible to find suitable baselines that can be easily tested.\n\nQUESTIONS:\nLine 216 states that \"each term in (3.2) can be made arbitrarily small by decreasing the amount of noise $\\sigma$ in the augmentation [...]\". Isn't the first term anti-proportional to the noise?\n\nLine 274: \"all sequences $s'\\_1 = s\\_1$ and  $s'\\_{h+1} = s\\_{h+1}$ [...]\". I'm wondering if this should be \"all sequences $s'\\_h = s\\_h$ and  $s'\\_{h+1} = s\\_{h+1}$ [...]\"!?\n\nLine 325: \"We let the expert policy $\\pi^*$ be the concatenation of policies $\\pi\\_h^*$ [...]\". Is this an additional assumption? Clearly, the expert will in general not apply this particular form of a hierarchical policy.\n\nLIMITATIONS:\nI think that the limitations are sufficiently clear.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-06-29"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper proposes a theoretical framework for studying the imitation of expert demonstrations in nonlinear dynamical systems. The framework guarantees accurate imitation of expert trajectories by invoking low-level controllers and a stochastic continuity property of the learned policy. The authors also introduce an algorithmic trick for ensuring total variation continuity, which combines a popular data-augmentation regimen with augmentation noise at execution time. The paper includes a detailed proof of the framework and empirical validation of the proposed augmentation strategy in simulated robotic manipulation tasks\n\nSTRENGTHS:\n1. The paper proposes a novel theoretical framework for studying the imitation of expert demonstrations in nonlinear dynamical systems.\n2. The framework guarantees accurate imitation of expert trajectories by invoking low-level controllers and a stochastic continuity property of the learned policy.\n3. The authors introduce a novel algorithmic trick for ensuring total variation continuity, which combines a popular data-augmentation regimen with augmentation noise at execution time.\n4. The paper includes a detailed proof of the framework and establishes stability guarantees for sequences of primitive controllers in non-linear control systems.\n5. The proposed augmentation strategy is empirically validated in simulated robotic manipulation tasks.\n6. The paper provides a comprehensive review of related work in the field of imitation learning and compares the proposed framework to past work in the appendix.\n\nWEAKNESSES:\n1. The paper relies on either synthesized primitive controllers or low-level stabilizing controllers built into problem environments to take advantage of local stability. Developing a more comprehensive approach to stability (perhaps one that does not require explicit gain synthesis and extends to non-smooth systems) is an exciting direction for future work.\n2. The paper focuses on simulated robotic manipulation tasks, and it is unclear how well the proposed framework would generalize to other domains or real-world scenarios.\n3. The paper does not provide a detailed analysis of the computational complexity of the proposed algorithmic trick for ensuring total variation continuity.\n4. The paper does not compare the proposed framework to state-of-the-art methods in the field of imitation learning on benchmark datasets.\n\nQUESTIONS:\n1. I think it is still important to comprehensive study of this method. What’s the influence of hyperparameters to this method? Should there be any ablation study included in this paper?\n2. What’s the computational complexity of the proposed algorithmic trick for ensuring total variation continuity?\n\nLIMITATIONS:\n1. The paper assumes that the low-level controllers are either learned or implicit in position-command control, which may not always be the case in practice.\n2. The paper does not provide a detailed analysis of the effect of hyperparameters on the performance of the proposed framework.\n3. The paper does not consider the effect of imperfect state estimation or perception on the accuracy of the imitation policy.\n4. The paper does not address the issue of exploration in the imitation learning setting, which is an important problem in reinforcement learning.", "TITLE": "Official Review", "RECOMMENDATION": 3, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-07-04"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper introduces an imitation learning framework for handling multimodal and non-Markovian demonstrations. This is done by using a hierarchy where rather than predicting low-level controls directly, the learner predicts sequences of primitive controllers. These controllers are time-varying affine mappings from states to control inputs. The authors show that this approach allows us to reformulate our task as imitation in a composite MDP, where composite states and actions refer to subsequences of trajectories and primitive controllers respectively. To handle multimodality, they use diffusion models (DDPM) to parameterize the learner. The primary contribution of the paper is to bound the loss of the proposed approach (TODA), relying primarily on the stability of the lower-level controllers. Some empirical results are demonstrated on some robot manipulation tasks, where the proposed algorithm outperforms the baseline policy.\n\nSTRENGTHS:\nThe high-level motivation for the approach is sound and the practical algorithm seems to perform well.\n\nThe paper is comprehensive and filled with technical content and derivations. I went through the analysis in the main body of the paper and did not find any glaring errors, although certain parts were hard to parse (see below).\n\nWEAKNESSES:\nHonestly I have a hard time following Section 4. More detailed questions are below. Given that the crux of the contribution lies in this theoretical analysis, I'm more than happy to adjust my score post-rebuttal if I can get some things cleared up.\n\nTraining a diffusion model to use low-level controllers as primitives as not new [1].\n\nWould have preferred a more comprehensive experimental evaluation with competitive baselines, but given the theoretical contribution I suppose this isn't fully necessary.\n\nThe proposed approach heavily relies on sufficiently smooth dynamics and the existence of an oracle which can synthesize primitive controllers given trajectories. This is fairly restrictive and prevents this work from being applicable to many existing control/robotics applications. Given that the contribution is primarily theoretical, this isn't a major issue for me.\n\n[1] Ajay, Anurag, et al. \"Is conditional generative modeling all you need for decision-making?.\" arXiv preprint arXiv:2211.15657 (2022).\n\nQUESTIONS:\nWhat is $d_{\\mathcal{A}}$? Not the formal definition, but conceptually what is it?\n\nWhat is the purpose of the replica policy? I'm not sure I understand how this part of the proof functions.\n\nIs the smoothing kernel $W_{\\sigma}$ the same as the data augmentation function?\n\nI don't understand L304-307.\n\nTypos:\n\nL52, \"... which is both both ...\"\n\nL60, \"Wassertstein-like\"\n\nL106, \"... we assume that (2.1) deterministic ...\"\n\nL342, \"subtley\"\n\n===\nI have read the rebuttal and bumped my score up slightly.\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 1, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-07-07"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis is a theoretical paper that studies imitation learning from demonstrations that exhibit multi-modality. This method separate the demonstrations into chunks, then uses diffusion models to learn a policy conditioned on the memory of hindsight prefix of the trajectory to output an affine time-varying controller. Gaussian noises are added to the demonstrations for data augmentation, and in the learned controllers' outputs to enhance performance.\n\nSTRENGTHS:\nOriginality: this paper is original in proposing an algorithm to learn hierarchical policy by leveraging diffusion models and injected noises from potentially non-markovian expert demonstrations.\nQuality: this paper has strong theoretical support. \nClarity: the algorithms and theories are clearly explained.\nSignificance: the paper provides a new insight to solving non-markovian imitation learning problems by combining trajectory optimization tools, distribution learning in a hierarhical manner.\n\nWEAKNESSES:\nIn the experiment, three benchmarks are utilized. However, it is not mentioned in the maintext nor in the appendix how non-markovian and multi-modal the demonstrations are. As a result, the validity of the proposed algorithm is less convincible despite all the theoretical guarantees. \nAlso, the symbols in this paper are too complicated to parse. There seem to be some symbols misusage in some of the formulas, making those formulas confusing. Please refer to the question section for the those potential errors.\n\nQUESTIONS:\n1. In line 5 of Algorithm 1, shouldn't the mean of \\gamma_h be 0 instead of a_h? In Eq.3.1, \\gamma_h is sampled from D to compose the a DDPM loss function of which the update rule matches Eq 2.2. Then, \\gamma_h corresponds to the \\gamma in line 155, which is a standard Gaussian.  \n2. If \\gamma_h indeed has its mean equal to a_h, how to understand the sum in Eq3.1? Why adding such noised action to e^{-\\alpha j} a_h? \n3. If \\gamma_h has its mean equal to 0, why matching s_\\theta,h(...) with \\gamma_h, a Gaussian noise? \n4. What do ...\\vee... \\leq max ... mean in line 275? Applying \\vee to real values?\n\nLIMITATIONS:\nThe author mentioned the limitation of the proposed approach in the paper. The limitation, however, does not undermine the quality and contribution of this paper.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-07"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis work proposes a theoretical framework for learning sequential decision-making with diffusion models. The authors show that with the conditions of (1) a suitable low-level \"stability\" guarantee and (2) a stochastic \"continuity\" property of the learned policy, the imitator can provably accurately estimate the \"non-Markovian\" policy of the demonstrators. The low-level \"stability\" is assumed in oracle primitive controllers, which are then composed by diffusion models that are learned with DDPM loss. Since the desired stochastic continuity is not guaranteed for an arbitrary demonstrator policy, the authors then propose a data augmentation technique that adds noise to the trajectories, making the resulting policies continuous. Their experiments validate that the proposed data augmentation alleviates imitator policy's degradation in noisy environments.\n\nSTRENGTHS:\nOriginality: I am not very familiar with the current literature of learning theory. But proving the learnability of a hierarchical imitator appears to be novel to me. \n\nQuality: The analysis presented in this paper is of high quality. I liked the formal narrative where each mathematical object is clearly defined before used. The authors also provide detailed derivation in their appendix. \n\nClarity: Even though this paper adopts heavy notation and formulas, the authors did a good job in explaining each mathematical object carefully. Through out my read, there are lines where I need to pause and think, but barely one makes me confused. The flow is very smooth. \n\nSignificance: I am not an expert in the literature of learning theory, so my evaluation of the significance may be biased. The proposed algorithm itself does not seem to be super significant. But conceptually connecting adding noises to \"total variation continuity\" is interesting to me.\n\nWEAKNESSES:\nMy major concern is that even though the presented analysis is interesting and self-contained, it doesn't seem to provide too much insight into developing a new algorithm or improving existing ones. The proposed algorithm is a consequence of the analysis, but it is not obvious why we need it. I hope the authors would like to discuss more about the potential influence of the presented analysis and algorithm in their feedback.\n\nQUESTIONS:\nApart from the question in Weakness, I hope the authors would like to answer the following: \n\nL172, what is K? It does not seem to be defined. \n\nI am not sure if this question is appropriate, but how can we know Assumptions 3.1, 3.2, and 3.3 are true in reality?\n\nLIMITATIONS:\nAs stated by the authors in L239-243, the major limitation of this work is that the desired \"stability\" is given since the proposed algorithm relies on either synthesized primitive controllers (in our analysis) or low-level stabilizing controllers built into problem environments (in our experiments). Relaxing such restriction can be a future research direction.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-19"}