SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Adam Block
Daniel Pfrommer
dpfrom@mit.edu
Max Simchowitz
msimchow@csail.mit.edu
Russ Tedrake
russt@mit.edu
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.
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
Figure 1: Consider demonstration trajectories exhibiting two modes: a "go left" and "go right" mode around an obstacle depicted in red and purple, respectively. To avoid compounding error, we imitate sequences of simple low-level feedback controllers we call "primitive controllers", not simply raw actions. Intuitively, primitive controllers provide "tubes" around each demonstration trajectory where the system can be stabilized. Depicted in yellow, our data-noising procedure described below "fills in the gaps" in the demonstration, switching between modes in a well-behaved manner, and whilst allowing the primitive controllers to manage the 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_{\mathrm{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  $\hat{\pi}$  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 \leq H$
$\mathbb{P}[\text{expert} \& \text{imitator trajectories disagree at some time } h \text{ by } \geq \varepsilon]$
$$
\leq \mathcal {O} _ {\mathrm {I S S}} \left(H \varepsilon + \frac {1}{\varepsilon^ {2}} \sum_ {h} \mathbb {E} _ {\text {e x p e r t}, h} \left[ \mathcal {W} _ {1} \left(\mathbb {P} _ {\text {e x p e r t a c t i o n s}}, \mathbb {P} _ {\text {i m i t a t o r a c t i o n s}}\right) \right]\right)
$$
where  $\mathbb{E}_{\mathrm{expert},h}[\mathcal{W}_1(\mathbb{P}_{\mathrm{expert~actions}},\mathbb{P}_{\mathrm{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  $\mathcal{O}_{\mathrm{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.
1. 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).
2. 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-command through inverse dynamics controllers [33, 2] or position-command controllers of end effectors [19] leads to substantially improved performance.
3. We also apply a subtle-yet-significant modification to a popular data noising strategy, which we show yields both theoretical and empirical benefits. 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 control-theoretic notions of stability, notably incremental stability, [68, 27, 50], which require continuity, Markovianity, and often determinism, and preclude the bifurcations permitted in our 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  $\mathbf{z}_{1:n} = (\mathbf{z}_1,\dots ,\mathbf{z}_n)$ . Norms  $\| \cdot \|$  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  $\mathcal{X}$  (which include real vector spaces), the space of Borel distributions on  $\mathcal{X}$  is denoted  $\Delta(\mathcal{X})$ . We rely heavily on couplings from optimal transport theory: given measures  $X \sim \mathsf{P}$  and  $X' \sim \mathsf{P}'$  on  $\mathcal{X}$  and  $\mathcal{X}'$  respectively,  $\mathcal{C}(\mathsf{P}, \mathsf{P}')$  denotes the space of joint distributions  $\mu \in \Delta(\mathcal{X} \times \mathcal{X}')$  called "couplings" such that  $(X, X') \sim \mu$  has marginals  $X \sim \mathsf{P}$  and  $X' \sim \mathsf{P}$ .  $\Delta(\mathcal{X} | \mathcal{Y})$  denotes the space of conditional probability distributions  $\mathsf{Q}: \mathcal{Y} \rightarrow \Delta(\mathcal{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  $\mathbf{I}_{\infty}(\mathcal{E})$  denotes the indicator taking value 1 if  $\mathcal{E}$  is true and  $\infty$  otherwise.
**Dynamics and Demonstrations.** We consider a discrete-time, control system with states  $\mathbf{x}_t \in \mathcal{X} := \mathbb{R}^{d_x}$ , and inputs  $\mathbf{u}_t \in \mathcal{U} := \mathbb{R}^{d_u}$ , obeying the following nonlinear dynamics
$$
\mathbf {x} _ {t + 1} = f \left(\mathbf {x} _ {t}, \mathbf {u} _ {t}\right), \quad t \geq 1. \tag {2.1}
$$
Given length  $T \in \mathbb{N}$ , we call sequences  $\rho_T = (\mathbf{x}_{1:T+1}, \mathbf{u}_{1:T}) \in \mathcal{P}_T := \mathcal{X}^{T+1} \times \mathcal{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 non-Markovian demonstrator, which allows for the multi-modality described in the Section 1.
Definition 2.1 (Expert Distribution). Let  $\mathcal{D}_{\mathrm{exp}} \in \Delta(\mathcal{P}_T)$  denote an expert distribution over trajectories to be imitated.  $\mathcal{D}_{\mathbf{x}_1}$  denotes the distribution of  $\mathbf{x}_1$  under  $\rho_T = (\mathbf{x}_{1:T+1}, \mathbf{u}_{1:T}) \sim \mathcal{D}_{\mathrm{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  $\mathcal{X} \to \mathcal{U}$  (redundantly) parameterized as  $\mathbf{x} \mapsto \bar{\mathbf{u}} + \bar{\mathbf{K}} (\mathbf{x} - \bar{\mathbf{x}})$ ; we call these primitive controllers, denoted with  $\kappa = (\bar{\mathbf{u}}, \bar{\mathbf{x}}, \bar{\mathbf{K}}) \in \mathcal{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  $\mathcal{D}_{\mathrm{exp}}$  may involve non-Markovian sequences of actions. We imitate these sequences via chunking policies. Fix a chunk length  $\tau_{\mathrm{chunk}} \in \mathbb{N}$  and observation length  $\tau_{\mathrm{obs}} \leq \tau_{\mathrm{chunk}}$ , and define time indices  $t_h = (h - 1)\tau_{\mathrm{chunk}} + 1$ . For simplicity, we assume  $\tau_{\mathrm{chunk}}$  divides  $T$ , and set  $H = T / \tau_{\mathrm{chunk}}$ . Given a  $\rho_T \in \mathcal{P}_T$ , define the trajectory-chunks and observation chunks
$$
\mathbf {s} _ {h} := \left(\mathbf {x} _ {t _ {h - 1}: t _ {h}}, \mathbf {u} _ {t _ {h - 1}: t _ {h} - 1}\right) \in \mathcal {S} := \mathscr {P} _ {\tau_ {\text {c h u n k}}} \quad (\text {t r a j e c t o r y - c h u n k s})
$$
$$
\mathcal {o} _ {h} := \left(\mathbf {x} _ {t _ {h} - \tau_ {\mathrm {o b s}} + 1: t _ {h}}, \mathbf {u} _ {t _ {h} - \tau_ {\mathrm {o b s}} + 1: t _ {h} - 1}\right) \in \mathcal {O} := \mathcal {P} _ {\tau_ {\mathrm {o b s}} - 1} \quad (\text {o b s e r v a t i o n - c h u n k s})
$$
for  $h > 1$ , and  $\mathbf{s}_1 = \mathbf{o}_1 = \mathbf{x}_1$  (for simplicity, we embed  $\mathfrak{o}_1$  into  $\mathcal{P}_{\tau_{\mathrm{obs}} - 1}$  via zero-padding). We call  $\tau_{\mathrm{chunk}}$ -length sequences of primitive controllers composite actions
$$
\mathsf {a} _ {h} = \kappa_ {t _ {h}: t _ {h - 1}} \in \mathcal {A} := \mathcal {K} ^ {\tau_ {\text {c h u n k}}} \quad (\text {c o m p o s i t e a c t i o n s})
$$
A chunking policy  $\pi = (\pi_h)$  consists of functions  $\pi_h$  mapping observation-chunks  $\mathsf{o}_h$  to distributions  $\Delta(\mathcal{A})$  over composite actions and interacting with the dynamics (2.1) by  $\mathsf{a}_h = \kappa_{t_h:t_{h-1}} \sim \pi_h(\mathsf{o}_h)$ , and executing  $\mathbf{u}_t = \kappa_t(\mathbf{x}_t)$ . We let  $d_{\mathcal{A}} = \tau_{\mathrm{chunk}}(d_x + d_u + d_x d_u)$  denote the dimension of the space  $\mathcal{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  $\varepsilon >0$  and a (chunking) policy  $\pi$ , the (marginal distribution) imitation loss is  $\mathcal{L}_{\mathrm{marg},\varepsilon}(\pi)\coloneqq \max_{t\in [T]}\inf_{\mu}\mathbb{P}_{\mu}\left[\max \left\{\| \mathbf{x}_{t + 1}^{\exp} - \mathbf{x}_{t + 1}^{\pi}\| ,\| \mathbf{u}_t^{\exp} - \mathbf{u}_t^\pi \| \right\} >\varepsilon \right]$ , where the infimum is over all couplings  $\mu$  between the distribution of  $\rho_T$  under  $\mathcal{D}_{\mathrm{exp}}$  and that induced by the policy  $\pi$  as described above, such that  $\mathbb{P}_{\mu}[\mathbf{x}_1^{\exp} = \mathbf{x}_1^\pi ] = 1$ .
Figure 2: Graphical comparison of an action-chunk based policy (left) as described in [19], versus the primitive-controller chunking policy (right) proposed in this paper. The primitive controller paradigm allows for stabilizing back to the original expert trajectory, whereas using generated actions in an open-loop fashion may cause divergence from the expert in the presence of unstable system dynamics. We refer to composite actions as the sequence of primitive controllers given on the right.
Under stronger conditions (whose necessity we establish), we can also imitate joint distributions over actions (Appendix J). Observe that  $\mathcal{L}_{\mathrm{fin},\varepsilon} \leq \mathcal{L}_{\mathrm{marg},\varepsilon}$ , 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,  $\mathcal{L}_{\mathrm{marg},\varepsilon}$  upper bounds the difference in expectation between any bounded and Lipschitz control cost decomposing across time steps, states and inputs, and  $\mathcal{L}_{\mathrm{fin},\varepsilon}$  upper bounds differences in final-state costs; see Appendix J for discussion.
Diffusion Models. Our analysis provides imitation guarantees when chunking policies  $\pi_h$  select  $\mathsf{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\in \Delta (\mathbb{R}^d)$  assuming that the score  $\nabla \log q$  is known to the learner. More precisely, suppose the learner is given an observation  $\mathsf{o}_h$  and wishes to sample  $\mathsf{a}_h\sim q(\cdot |\mathsf{o}_h)$  for some family of probability kernels  $q(\cdot |\cdot)$ . A DDPM starts with some  $\mathsf{a}_h^0$  sampled from a standard Gaussian noise and iteratively "denoises" for each DDPM-time step  $0\le j < J$ :
$$
\mathsf {a} _ {h} ^ {j} = \mathsf {a} _ {h} ^ {j - 1} - \alpha \cdot \mathbf {s} _ {\theta , h} \left(\mathsf {a} _ {h} ^ {j - 1}, \mathsf {o} _ {h}, j\right) + 2 \cdot \mathcal {N} \left(0, \alpha^ {2} \mathbf {I}\right), \tag {2.2}
$$
where  $\mathbf{s}_{\theta,h}(\mathsf{a}_h^j, \mathsf{o}_h, j)$  estimates the true score  $\mathbf{s}_{\star,h}(\mathsf{a}_h, \mathsf{o}_h, \alpha j)$ , formally defined for any continuous argument  $t \leq J\alpha$  to be  $\mathbf{s}_{\star,h}(\mathsf{a}, \mathsf{o}_h, t) := \nabla_{\mathsf{a}} \log q_{h,[t]}^{\star}(\mathsf{a} \mid \mathsf{o}_h)$ , where  $q_{h,[t]}^{\star}(\cdot | \mathsf{o}_h)$  is the distribution of  $e^{-t}\mathsf{a}_h^{(0)} + \sqrt{1 - e^{-2t}}\gamma$  with  $\mathsf{a}_h^{(0)}$  is sampled from the target distribution we which to sample from, and  $\gamma \sim \mathcal{N}(0, \mathbf{I})$  is a standard Gaussian. We denote by  $\mathrm{DDPM}(\mathbf{s}_{\theta}, \mathsf{o}_h)$  the law of  $\mathsf{a}_h^J$  sampled according to the DDPM using  $\mathbf{s}_{\theta}(\cdot, \mathsf{o}_h, \cdot)$  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)  $\kappa_{1:T} \in \mathcal{K}^T$  is consistent with a trajectory  $\pmb{\rho} = (\mathbf{x}_{1:T+1}, \mathbf{u}_{1:T}) \in \mathcal{P}_T$  if  $\bar{\mathbf{x}}_t = \mathbf{x}_t$  and  $\bar{\mathbf{u}}_t = \mathbf{u}_t$  for all  $t \in [T]$ ; note that this implies that  $\kappa_t(\mathbf{x}_t) = \mathbf{u}_t$  for all  $t$ . A synthesis oracle synth maps  $\mathcal{P}_T \to \mathcal{K}^T$  such that, for all  $\pmb{\rho}_T \in \mathcal{P}_T$ ,  $\kappa_{1:T} = \mathrm{synth}(\pmb{\rho}_T)$  is consistent with  $\pmb{\rho}_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  $\mathsf{a} = (\bar{\mathbf{u}}_{1:\tau_{\mathrm{chunk}}},\bar{\mathbf{x}}_{1:\tau_{\mathrm{chunk}}},\bar{\mathbf{K}}_{1:\tau_{\mathrm{chunk}}} )$  and  $\mathsf{a}' = (\bar{\mathbf{u}}_{1:\tau_{\mathrm{chunk}}}^{\prime},\bar{\mathbf{x}}_{1:\tau_{\mathrm{chunk}}}^{\prime},\bar{\mathbf{K}}_{1:\tau_{\mathrm{chunk}}}^{\prime})$ , we define
$$
\mathrm {d} _ {\max } \left(\mathsf {a}, \mathsf {a} ^ {\prime}\right) := \max  _ {1 \leq k \leq \tau_ {\text {c h u n k}}} \left(\| \bar {\mathbf {u}} _ {k} - \bar {\mathbf {u}} _ {k} ^ {\prime} \| + \| \bar {\mathbf {x}} _ {k} - \bar {\mathbf {x}} _ {k} ^ {\prime} \| + \| \bar {\mathbf {K}} _ {k} - \bar {\mathbf {K}} _ {k} ^ {\prime} \|\right). \tag {3.1}
$$