SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Distances between policies are defined via natural optimal transport costs. Given two policies  $\pi = (\pi_h), \pi' = (\pi_h')$  and observation chunk  $o_h$ , we define an induced optimal transport cost
$$
\Delta_ {\varepsilon} \left(\pi_ {h} \left(\mathsf {o} _ {h}\right), \pi_ {h} ^ {\prime} \left(\mathsf {o} _ {h}\right)\right) := \inf  _ {\mu} \mathbb {P} _ {\left(\mathsf {a} _ {h}, \mathsf {a} _ {h} ^ {\prime}\right) \sim \mu} \left[ \mathrm {d} _ {\max } \left(\mathsf {a} _ {h}, \mathsf {a} _ {h} ^ {\prime}\right) > \varepsilon \right],
$$
where the  $\inf_{\mu}$  denotes the infimum over all couplings between  $\mathsf{a}_h\sim \pi_h(\mathsf{o}_h)$  and  $\mathsf{a}_h^\prime \sim \pi_h^\prime (\mathsf{o}_h)$ .  $\Delta_{\varepsilon}$  corresponds to a relaxed Lévy-Prokhorov metric [64], and can always be bounded, via Markov's inequality, by  $\Delta_{\varepsilon}(\pi_h,\pi_h^{\prime}\mid \mathsf{o}_h)\leq \frac{1}{\varepsilon}\mathcal{W}_{1,\mathrm{d}_{\max}}(\pi_h(\mathsf{o}_h),\pi_h^{\prime}(\mathsf{o}_h))$  where  $\mathcal{W}_{1,\mathrm{d}_{\max}}(\pi_h(\mathsf{o}_h),\pi_h^{\prime}(\mathsf{o}_h))$  denotes the 1-Wasserstein distance between  $\mathsf{a}_h\sim \pi_h(\mathsf{o}_h)$  and  $\mathsf{a}_h^\prime \sim \pi_h^\prime (\mathsf{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  $\kappa : \mathbb{R}^{d_x} \to \mathbb{R}^{d_u}$ , define the closed loop dynamic map  $f_{\mathrm{cl},\kappa}(\mathbf{x},\delta \mathbf{u}) \coloneqq f(\mathbf{x},\kappa (\mathbf{x}) + \delta \mathbf{u})$ . Thus, composite action  $a$  is consistent with a trajectory chunk  $s = (x_{1:\tau_{\mathrm{chunk}} + 1}, u_{1:\tau})$  if  $x_{t + 1} = f_{\mathrm{cl},\kappa_t}(x,0)$  for  $1 \leq t \leq \tau_{\mathrm{chunk}}$ .
Definition 3.1 (Time-Varying Incremental Stability). Let  $\gamma (\cdot)$  be a class  $\mathsf{K}$  function,  $\beta (\cdot ,\cdot)$  be class KL function, and let  $\mathsf{a} = (\kappa_1,\kappa_2,\dots ,\kappa_\tau)$  denote a sequence of primitive controllers (i.e. a composite action when  $\tau = \tau_{\mathrm{chunk}}$  ). Given a sequence of input perturbations  $\delta \mathbf{u}_{1:\tau}\in (\mathbb{R}^{d_u})^\tau$  and initial condition  $\pmb {\xi}\in \mathbb{R}^{d_x}$ , let  $\mathbf{x}_{i + 1}^{\mathsf{a}}(\delta \mathbf{u}_{1:\tau},\pmb {\xi}) = f_{\mathrm{cl},\kappa_i}(\mathbf{x}_i^{\mathsf{a}}(\delta \mathbf{u}_{1:\tau},\pmb {\xi}),\delta \mathbf{u}_i)$ , with  $\mathbf{x}_1^{\mathsf{a}} = \pmb{\xi}$ . We say that composite action  $a$  is time-varying incrementally input-to-state stable (t-Iss) with moduli  $\gamma (\cdot),\beta (\cdot ,\cdot)$  if for all  $\pmb {\xi},\pmb{\xi}'\in \mathbb{R}^{d_x},0\leq i\leq \tau$ $\| \mathbf{x}_i^{\mathsf{a}}(\mathbf{0}_{1:\tau},\pmb {\xi}) - \mathbf{x}_i^{\mathsf{a}}(\delta \mathbf{u}_{1:\tau},\pmb{\xi}')\| \leq \beta (\| \pmb {\xi} - \pmb{\xi}'\| ,\tau) +$ $\gamma$  ( $\max_{1\leq s\leq i - 1}\| \delta \mathbf{u}_s\|$ ). Given parameters  $c_{\gamma},c_{\xi} > 0$  we say that a is local-t-Iss at  $\pmb{\xi}_0$  if the above holds only for all  $\pmb {\xi},\pmb{\xi}',\delta \mathbf{u}_{1:\tau}$  such that  $\| \pmb {\xi} - \pmb{\xi}_0\| ,\| \pmb{\xi}' - \pmb{\xi}_0\| \leq c_{\xi}$  and  $\max_t\| \delta \mathbf{u}_t\| \leq c_{\gamma}$ .
Incremental stability implies that as the initial conditions  $\| \pmb {\xi} - \pmb{\xi}'\| \to 0$  and  $\max_{0\leq s\leq i - 1}\| \delta \mathbf{u}_t\| \to 0$  the trajectories induced by taking rolling out a from  $\pmb{\xi}$ , and rolling out a from  $\pmb{\xi}'$  with additive input perturbations  $\delta \mathbf{u}_{1:\tau}$  tend to zero in norm. This behavior needs only hold for initial conditions in a small neighborhood of a nominal state  $\pmb{\xi}_0$ . Importantly, the perturbations  $\delta \mathbf{u}_{1:\tau}$  are fixed perturbations of inputs, applied to the closed loop behavior under the controllers. Our notion of incremental stability are similar too, but subtly 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  $\rho_T = (\mathbf{x}_{1:T+1}, \mathbf{u}_{1:T}) \sim \mathcal{D}_{\mathrm{exp}}$ , and let  $\kappa_{1:T} = \mathrm{synth}(\rho_T)$ , partitioned into composite actions  $a_{1:H}$ , with  $\kappa_t(\mathbf{x}) = \bar{\mathbf{K}}_t(\mathbf{x} - \bar{\mathbf{x}}_t) + \bar{\mathbf{u}}_t$ . We assume that, with probability one,  $\kappa_{1:T}$  is consistent with  $\rho_T^2$ , and that, for each  $1 \leq h \leq H$ ,  $a_h = (\kappa_{t_h:t_h + \tau_{\mathrm{chunk}} - 1})$  is local t-ISS at  $\mathbf{x}_{t_h}$  with moduli  $\gamma, \beta$  and parameters  $c_\beta, c_\xi > 0$ . We further assume that  $\gamma$  and  $\beta$  take the form
$$
\gamma (u) = \bar {c} _ {\gamma} \cdot u, \quad \beta (u, k) = \bar {c} _ {\beta} e ^ {- (k - 1) L _ {\beta}} \cdot u, \quad \bar {c} _ {\gamma}, \bar {c} _ {\beta} > 0, \quad L _ {\beta} \in (0, 1 ].
$$
<sup>1</sup>Below, we recall definitions of classes of comparison functions in nonlinear control [38] as follows: we say a univariate function  $\gamma : \mathbb{R}_{\geq 0} \to \mathbb{R}_{\geq 0}$  is class K if it is strictly increasing and satisfies  $\gamma(0) = 0$ . We say a bivariate function  $\beta : \mathbb{R}_{\geq 0} \times \mathbb{Z}_{\geq 0} \to \mathbb{R}_{\geq 0}$  is class KL if  $x \mapsto \beta(x, t)$  is class K for each  $t \geq 0$ , and  $t \mapsto \beta(x, t)$  is nonincreasing in  $t$ .
Note that this implies  $\bar{\mathbf{x}}_t = \mathbf{x}_t$  and  $\bar{\mathbf{u}}_t = \mathbf{u}_t$ .
Figure 3: Graphical representation of total variation continuity (TVC) using the running "left mode/right mode" example. Panel  $(a)$  depicts a policy  $\pi$  which is TVC, and thus interpolates between left and right modes probabilistically. Importantly, the TVC property applies to the distribution over composite actions, i.e. sequences of primitive controllers, as, in Definition 3.3; this ensures, for example, that following the left mode from slightly to the right of the obstacle (purple dotted line) still stabilizes to the idealized left mode trajectory (red). In panel  $(b)$ , we consider a policy which for TVC applies to the sequences of raw control inputs (which is not what occurs in Definition 3.3). This can lead to naive mode-switching that collides with the gray obstacle.
Lastly, we assume that for the expert trajectories and the primitive controllers drawn as above, it holds that satisfy  $\max \{\| \mathbf{x}_t\| ,\| \mathbf{u}_t\| \} \leq R_{\mathrm{dyn}}$  and  $\| \mathbf{K}_t\| \leq R_{\mathbf{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  $\hat{\pi}$  that can compute the conditional distribution of composite actions at time steps given observation-chunks, then a stochastically smoothed version of this policy,  $\hat{\pi}_{\sigma}$ , has low imitation error. Define, for any length  $\tau \in \mathbb{N}$ , the trajectory distance between trajectories  $\pmb{\rho} = (\mathbf{x}_{1:\tau + 1}, \mathbf{u}_{1:\tau})$ ,  $\pmb{\rho}' = (\mathbf{x}_{1:\tau + 1}', \mathbf{u}_{1:\tau}') \in \mathcal{P}_{\tau}$
$$
\mathrm {d} _ {\text {t r a j}} \left(\boldsymbol {\rho}, \boldsymbol {\rho} ^ {\prime}\right) := \max  _ {1 \leq k \leq \tau + 1} \| \mathbf {x} _ {k} - \mathbf {x} _ {k} ^ {\prime} \| \vee \max  _ {1 \leq k \leq \tau} \| \mathbf {u} _ {k} - \mathbf {u} _ {k} ^ {\prime} \|. \tag {3.2}
$$
In particular, we define  $\mathsf{d}_{\mathrm{traj}}(\mathsf{o}_h,\mathsf{o}_h^{\prime})$  and  $\mathsf{d}_{\mathrm{traj}}(\mathsf{s}_h,\mathsf{s}_h^{\prime})$  by viewing these as trajectories of length  $\tau_{\mathrm{obs}} - 1$  and  $\tau_{\mathrm{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  $\pi = (\pi_h)$  is total variation continuous with modulus  $\gamma_{\mathrm{TVC}}: \mathbb{R}_{\geq 0} \to \mathbb{R}_{\geq 0}$ , written  $\gamma_{\mathrm{TVC}}$ -TVC, if, for all  $h \in [H]$  and any observation-chunks  $\mathsf{o}_h, \mathsf{o}_h' \in \mathcal{P}_{\tau_{\mathrm{obs}} - 1}$ ,  $\overline{\mathrm{TV}}(\pi_h(\mathsf{o}_h), \pi_h(\mathsf{o}_h')) \leq \gamma_{\mathrm{TVC}}(\mathsf{d}_{\mathrm{traj}}(\mathsf{o}_h, \mathsf{o}_h'))$ .
We depict the TVC property using our running left-right obstacle example in Figure 3. We stress that, in Definition 3.3, the TV bound on  $\mathsf{TV}(\pi_h(\mathsf{o}_h),\pi_h(\mathsf{o}_h'))$  applies to the composite actions consisting of primitive controllers  $\mathsf{a}_h = \kappa_{t_h:t_h + \tau_{\mathrm{chunk}} - 1}\sim \pi_h(\mathsf{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 \in [H]$ , we let  $\mathcal{D}_{\mathrm{exp},h}$  denote the joint distribution of  $(\mathsf{a}_h, \mathsf{o}_h)$ , induced by drawing a trajectory  $\pmb{\rho}_T = (\mathbf{x}_{1:T+1}, \mathbf{u}_{1:T}) \sim \mathcal{D}_{\mathrm{exp}}$  from the expert distribution,  $\kappa_{1:T} = \mathrm{synth}(\pmb{\rho}_T)$  be the associated primitive controllers, letting  $\mathsf{o}_h = (\mathbf{x}_{t_h - \tau_{\mathrm{obs}} + 1:t_h}, \mathbf{u}_{t_h - \tau_{\mathrm{obs}} + 1:t_h - 1})$  be the associated observation-chunk at time  $h$ , and  $\mathsf{a}_h = \kappa_{t_h:t_{h+1}-1}$  the associated composite action. We let  $\pi_h^*(\cdot): \mathcal{O} \to \Delta(\mathcal{A})$  denote the condition distribution of  $\mathsf{a}_h \mid \mathsf{o}_h$  under  $\mathcal{D}_{\mathrm{exp},h}$ .
The conditional distributions  $\pi_h^\star (\cdot)$  are estimated when training a generative model to predict  $\mathsf{a}_h$  from observations  $\mathsf{o}_h$ . Note that  $\pi_h^\star (\cdot)$  (and  $\mathcal{D}_{\mathrm{exp},h}$ ) is defined in terms of both expert demonstration
from  $\mathcal{D}_{\mathrm{exp}}$  and the associated synthesized primitive controllers. In Lemma J.6, we show that when the synthesis oracle  $\kappa_{1:T} = \mathrm{synth}(\pmb{\rho}_T)$  produces primitive controllers consistent with the trajectories, then  $\pi^{\star} = (\pi_h^{\star})$  produces the same marginals over states as  $\mathcal{D}_{\mathrm{exp}}$ ; that is,  $\mathcal{L}_{\mathrm{marg},\varepsilon}(\pi^{\star}) = 0$ .
Theorem 1. Suppose Assumption 3.1 holds, and suppose that  $0 \leq \varepsilon < c_{2}$ , and  $\tau_{\mathrm{chunk}} \geq c_{3}$ . Then, for any non-decreasing non-negative  $\gamma_{\mathrm{TVC}}(\cdot)$  and  $\gamma_{\mathrm{TVC}}$ -TVC chunking policy  $\hat{\pi}$ , it holds that  $\mathcal{L}_{\mathrm{marg},\varepsilon}(\hat{\pi}) \leq H\gamma_{\mathrm{TVC}}(\varepsilon) + \sum_{h=1}^{H} \mathbb{E}_{\mathfrak{o}_h \sim \mathcal{D}_{\mathrm{exp},h}} \Delta_{(\varepsilon / c_1)}(\pi_h^*(\mathfrak{o}_h), \hat{\pi}_h(\mathfrak{o}_h))$ , which is at most  $H\gamma_{\mathrm{TVC}}(\varepsilon) + \frac{c_1}{\varepsilon} \sum_{h=1}^{H} \mathbb{E}_{\mathfrak{o}_h \sim \mathcal{D}_{\mathrm{exp},h}}[\mathcal{W}_{1,\mathbf{d}_{\max}}(\pi_h^*(\mathfrak{o}_h), \hat{\pi}_h(\mathfrak{o}_h))]$ .
The above result reduces the marginal imitation error of  $\hat{\pi}$  to the sum over optimal transport errors between  $\hat{\pi}$  and a  $|\mathsf{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  $\mathsf{a} \mid \tilde{\mathfrak{o}}_h$ , where  $\tilde{\mathfrak{o}}_h \sim \mathcal{N}(\mathfrak{o}_h, \sigma^2\mathbf{I})$  adds  $\sigma^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  $\hat{\pi} = (\hat{\pi}_h)$  be a chunking policy. We define the smoothed policy  $\hat{\pi}_{\sigma} = (\hat{\pi}_{\sigma,h})$  by letting  $\hat{\pi}_{\sigma,h}(\cdot|\mathfrak{o}_h)$  be distributed as  $\hat{\pi}_h(\cdot|\tilde{\mathfrak{o}}_h)$ , where  $\tilde{\mathfrak{o}}_h \sim \mathcal{N}(\mathfrak{o}_h, \sigma^2\mathbf{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  $\mathcal{D}_{\mathrm{exp},h}$  be as in Definition 3.4. Define  $\mathcal{D}_{\mathrm{exp},\sigma,h}$  as the distribution over  $(\tilde{\mathsf{o}}_h,\mathsf{a}_h)$  generated by  $(\mathsf{o}_h,\mathsf{a}_h)\sim \mathcal{D}_{\mathrm{exp},h}$  and  $\tilde{\mathfrak{o}}_h\sim \mathcal{N}(\mathfrak{o}_h,\sigma^2\mathbf{I})$ . We define the deconvolution policy  $\pi_{\mathrm{dec},\sigma,h}^{\star}(\tilde{\mathsf{o}}_h)$  as the conditional distribution of  $\mathsf{a}_h$  under  $\mathcal{D}_{\mathrm{exp},\sigma,h}$ .
Analogously to  $\pi^{\star}$ , the policy  $\pi_{\mathrm{dec},\sigma,h}^{\star}$  is what a generative model trained to generate  $\mathsf{a}_h$  from noisy observations  $\tilde{\mathsf{o}}_h$  of  $\mathsf{o}_h \sim \mathcal{D}_{\mathrm{exp}}$  learns to generate. Our next theorem states that, if our  $\hat{\pi}$  approximates the idealized conditional distribution of composite actions given noised observations, then  $\hat{\pi}_{\sigma}$ , 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},\ldots ,c_{5} > 0$  defined in Definition 3.2, and let  $\Theta_{\mathrm{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  $\varepsilon \leq \Theta_{\mathrm{iss}}(1)$ , if we choose  $\sigma = \varepsilon /\Theta_{\mathrm{iss}}(\sqrt{d_x} +\log (1 / \varepsilon))$  and let  $\tau_{\mathrm{chunk}}\leq c_3$  and  $\tau_{\mathrm{chunk}} - \tau_{\mathrm{obs}}\geq \frac{1}{L_{\beta}}\log (c_1 / \varepsilon)$
$$
\mathcal {L} _ {\mathrm {m a r g}, \varepsilon} (\hat {\pi} _ {\sigma}) \leq \Theta_ {\mathrm {I s s}} \left( \right.\varepsilon H \sqrt {\tau_ {\mathrm {o b s}}} \cdot \left(\sqrt {d _ {x}} + \log \left(\frac {1}{\varepsilon}\right)\right) + \sum_ {h = 1} ^ {H} \mathbb {E} _ {\tilde {\mathbf {o}} _ {h} \sim \mathcal {D} _ {\exp , \sigma , h}} \left[ \Delta_ {(\varepsilon^ {2})} \left(\pi_ {\mathrm {d e c}, \sigma , h} ^ {\star} \left(\tilde {\mathbf {o}} _ {h}\right), \hat {\pi} _ {h} \left(\tilde {\mathbf {o}} _ {h}\right)\right)\right],
$$
which is upper bounded by at most  $\Theta_{\mathrm{Iss}}\left(\varepsilon H\sqrt{\tau_{\mathrm{obs}}}\cdot \left(\sqrt{d_x} +\log (1 / \varepsilon)\right)\right.$ $+\frac{1}{\varepsilon^2}\sum_{h = 1}^{H}\mathbb{E}_{\mathfrak{o}_h\sim \mathcal{D}_{\mathrm{exp},\sigma ,h}}\left[\mathcal{W}_{1,\mathbf{d}_{\max}}\left(\pi_{\mathrm{dec},\sigma ,h}^\star (\tilde{\mathfrak{o}}_h),\hat{\pi}_h(\tilde{\mathfrak{o}}_h)\right)\right].$
To reiterate, Theorem 2 guarantees imitation of the distribution of marginals and final states of  $\mathcal{D}_{\mathrm{exp}}$  by replacing the explicit TVC assumption with noising, and the resulting guarantee applies to the smoothed policy  $\hat{\pi}_{\sigma}$  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
Figure 4: Multi-modal demonstrations traverse an obstacle left or right, exhibiting a pure bifurcation. We consider perturbing expert data on the right mode (blue circle) to a noised datum (gray circle). We show that generative behavior clones learn to deconvolve this noise, creating a virtual "replica" sample (red circle) following the left mode, such that the replica and original are i.i.d. given the noised one. When the red circle's primitive controllers are rolled from the blue circle, this leads to a trajectory (yellow circle) which interpolates across the bifurcations. Marginalizing over this process, the yellow trajectories probabilistically interpolate between red and blue modes, and (approximately) match the per-time-step marginal over expert distributions.
4 HINT: Instantiating Data Noising with DDPMs
We now instantiate Theorem 2 by showing that one can learn a policy  $\hat{\pi}$  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 Hierarchical Imitation via Noising at Inference Time (HINT)
[h] 1: Initialize Synthesis oracle synth, sample sizes  $N_{\mathrm{exp}}$ $N_{\mathrm{aug}}\in \mathbb{N},\sigma \geq 0$  ,DDPM step size  $\alpha >0$  DDPM horizon  $J$  , function class  $\{\mathbf{s}_{\theta}\}_{\theta \in \Theta}$  , gain magnitude  $R > 0$  , empty data buffer  $\mathfrak{D}\gets \emptyset$  .  $\%$  For no smoothing, set  $\sigma = 0$  and  $N_{\mathrm{aug}} = 1$    
2: for  $n = 1,2,\ldots N_{\mathrm{exp}}$  do   
3: Sample  $\rho_T = (x_{1:T + 1},u_{1:T})\sim \mathcal{D}_{\mathrm{exp}}$  and set  $\kappa_{1:T} = \mathrm{synth}(\pmb {\rho})$ $\%$  Segment  $\mathsf{o}_{1:H}$  from  $\rho_T$  and a1:H from  $\kappa_{1:T}$    
4: for  $i = 1,2,\dots ,N_{\mathrm{aug}}$  and  $h = 1,2,\dots ,H$  do   
5: Sample  $\tilde{\mathfrak{o}}_h\sim \mathcal{N}(\mathfrak{o}_h,\sigma^2\mathbf{I}),j_h\sim \mathrm{Unif}([J])$  and  $\gamma_{h}\sim \mathcal{N}(0,(j_{h}\alpha)^{2}\mathbf{I}).$    
6:  $\mathfrak{D}\gets \mathfrak{D}.$  append  $(\{(a_h,\tilde{o}_h,j_h,\gamma_h,h)\})$    
7: Fit  $\theta \in \arg \min_{\theta \in \Theta}\mathcal{L}_{\mathrm{DDPM}}(\theta ,\mathfrak{D})$  , and let  $\hat{\pi} = (\hat{\pi}_h)$  be given by  $\hat{\pi} (\cdot |\mathfrak{o}_h) = \mathrm{DDPM}(\mathbf{s}_{\theta ,h},\mathfrak{o}_h)$    
8: return  $\hat{\pi}_{\sigma} = (\hat{\pi}_{\sigma ,h})$  , by smoothing  $\hat{\pi}$  as per Definition 3.5.
rithm, HINT (Algorithm 1) combines DDPM-learning of chunked policies as in [19] with a popular form of data-augmentation [36]. We collect  $N_{\mathrm{exp}}$  expert trajectories, synthesize gains, and segment trajectories into observation-chunks  $\mathfrak{o}_h$  and composite actions  $\mathfrak{a}_h$  as described in Section 2. We perturb each  $\mathfrak{o}_h$  to form  $N_{\mathrm{aug}}$  chunks  $\tilde{\mathfrak{o}}_h$ , as well as horizon indices  $j \in [J]$  and inference noises  $\gamma \sim \mathcal{N}(0, (\alpha j_h)^2\mathbf{I})$ , and add these tuples  $(\mathfrak{a}_h, \tilde{\mathfrak{o}}_h, j_h, \gamma_h, h)$  to our data  $\mathfrak{D}$ . We end the training phase by minimizing the standard DDPM loss [62]:
$$
\mathcal {L} _ {\mathrm {D D P M}} (\theta , \mathfrak {D}) = \sum_ {\left(\mathfrak {a} _ {h}, \bar {\mathfrak {s}} _ {h}, j _ {h}, \gamma_ {h}, h\right) \in \mathfrak {D}} \left| \left| \boldsymbol {\gamma} _ {h} - \mathbf {s} _ {\theta , h} \left(e ^ {- \alpha j} \mathfrak {a} _ {h} + \sqrt {1 - e ^ {- 2 \alpha j}} \gamma_ {h}, \bar {\mathfrak {o}} _ {h}, j _ {h}\right) \right| \right| ^ {2}. \tag {4.1}
$$
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  $\mathrm{DDPM}(\mathbf{s}_{\theta,h},\cdot)\circ \mathcal{N}(\mathfrak{o}_h,\sigma^2\mathbf{I})$  means, given  $\mathfrak{o}_h$ , we perturb it to  $\tilde{\mathfrak{o}}_h\sim \mathcal{N}(\mathfrak{o}_h,\sigma^2\mathbf{I})$ , and sample  $\mathfrak{a}_h\sim \mathrm{DDPM}(\mathbf{s}_{\theta,h},\tilde{\mathfrak{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_{\mathrm{exp}}$ , polynomially large in problem parameters,  $\mathcal{L}_{\mathrm{marg},\varepsilon}(\widehat{\pi}_{\sigma}) \leq \Theta \left( \varepsilon H \sqrt{\tau_{\mathrm{obs}}} \left( \sqrt{d_x} + \log(1 / \varepsilon) \right) \right)$  with high probability.
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
Figure 5: Performance of diffusing chunks of actions  $\bar{\mathbf{u}}_{t_{h - 1}:t_h - 1}$  ("No Gains") versus jointly diffusing actions  $u_{t_{h - 1}:t_h - 1}$ , reference states  $\bar{\mathbf{x}}_{t_{h - 1}:t_h - 1}$  and gains  $\mathbf{K}_{t_{h - 1}:t_h - 1}$  for a 2-D quadrotor system with thrust-and-torque-based control. Different noise levels  $\sigma$  and number of trajectories  $N$  are shown. Mean and standard deviation are shown across 5 training seeds.
Figure 6: Performance of baseline  $\hat{\pi}$  and noise-injected  $\hat{\pi} \circ W_{\sigma}$  HINT policy for different  $\sigma$ . We use 4 training seeds with 50 and 22 test trajectories per seed for PushT and Can and Square Environments respectively. Mean and standard deviation of the test performance on the 3 best checkpoints across the 4 seeds are plotted. The  $\sigma$  values correspond to noise in the normalized  $[-1,1]$  range.
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.