CoolFace
Apppublic

lerobot/robot-learning-tutorial

sourceHugging Faceupdated 1y agoView on Hugging Face
508likes
04_imitation_learning.mdx468 linesDownload Raw Back to chapters
1import ResponsiveImage from '../../components/ResponsiveImage.astro'2import Ch4BcTrajectories from '../assets/image/ch4/ch4-bc-trajectories.png'3import Ch4ObservationActionMapping from '../assets/image/ch4/ch4-observation-action-mapping.png'4import Ch4IssuesWithBc from '../assets/image/ch4/ch4-issues-with-bc.png'5import Ch4TaskEffectOnPairs from '../assets/image/ch4/ch4-task-effect-on-pairs.png'6import Ch4LatentVariableModel from '../assets/image/ch4/ch4-latent-variable-model.png'7import Ch4ManyLatents from '../assets/image/ch4/ch4-many-latents.png'8import Ch4DiffusionRobotActions from '../assets/image/ch4/ch4-diffusion-robot-actions.png'9import Ch4ActionVsObservationDistribution from '../assets/image/ch4/ch4-action-vs-observation-distribution.png'10import Ch4NormalizingFlows from '../assets/image/ch4/ch4-normalizing-flows.png'11import Ch4DiffusionVsFlowmatching from '../assets/image/ch4/ch4-diffusion-vs-flowmatching.png'12import Ch4Act from '../assets/image/ch4/ch4-act.png'13import Ch4ActEncoder from '../assets/image/ch4/ch4-act-encoder.png'14import Ch4ActDecoder from '../assets/image/ch4/ch4-act-decoder.png'15import Ch4DiffusionPolicy from '../assets/image/ch4/ch4-diffusion-policy.png'16import Ch4AsyncInference from '../assets/image/ch4/ch4-async-inference.png'17import Ch4Queues from '../assets/image/ch4/ch4-queues.png'18 19 20# Robot (Imitation) Learning21 22::: epigraph23*The best material model for a cat is another, or preferably the same cat*24 25Norbert Wiener26:::27 28> **TL;DR**29> Behavioral Cloning provides a natural platform to learn from real-world interactions without the need to design any reward function, and generative models prove more effective than point-wise policies at dealing with multimodal demonstration datasets.30 31<ResponsiveImage src={Ch4BcTrajectories} alt="(A) Average (with standard deviation) evolution of the actuation levels over the first 5 recorded episodes in /svla_so101_pickplace. Proprioperceptive state provide invaluable to determine the robot's state during an episode. (B) Camera frames are also recorded alongside measurements on the robot's state, capturing information about the robot's interaction with its environment." id="fig-fig:ch4-bc-trajectories" />32 33*(A) Average (with standard deviation) evolution of the actuation levels over the first 5 recorded episodes in /svla_so101_pickplace. Proprioperceptive state provide invaluable to determine the robot's state during an episode. (B) Camera frames are also recorded alongside measurements on the robot's state, capturing information about the robot's interaction with its environment.*34 35Learning from human demonstrations provides a pragmatic alternative to the reinforcement-learning pipeline discussed in the following section.36Indeed, in real-world robotics online exploration is typically , and designing (dense) reward signals is a process.37In general, success detection itself may often require bespoke instrumentation, while episodic training demands reliable resets---all factors complicating training RL algorithms on hardware at scale.38Behavioral Cloning (BC) sidesteps these constraints by casting control an imitation learning problem, leveraging previously collected expert demonstrations.39Most notably, by learning to imitate autonomous systems naturally adhere to the objectives, preferences, and success criteria implicitly encoded in the data, which obviates reduces early-stage exploratory failures and obviates hand-crafted reward shaping altogether.40 41Formally, let $\mathcal D = \{ \tau^\{(i)\} \}_\{i=1\}^N$ be a set of expert trajectories, with $\tau^\{(i)\} = \{(o_t^\{(i)\}, a_t^\{(i)\})\}_\{t=0\}^\{T_i\}$ representing the $i$-th trajectory in $\mathcal D $, $ o_t \in \mathcal\{O\}$ denoting observations (e.g., images and proprioception altogether), and $a_t \in \mathcal\{A\}$ the expert actions.42Typically, observations $o \in \mathcal\{O\}$ consist of both image and proprioperceptive information, while actions $a \in \mathcal\{A\}$ represent control specifications for the robot to execute, e.g. a joint configuration.43Note that differently from the following section, in the imitation learning context $\mathcal D $ denotes an offline dataset collecting $ N$ length-$T_i$ reward-free (expert) human trajectories $\tau^\{(i)\}$, and *not* the environment dynamics.44Similarily, in this section $\tau^\{(i)\}$ represent a length-$T_i$ trajectory of observation-action pairs, which crucially *omits entirely any reward* information.45Figure the referenced figure graphically shows trajectories in terms of the average evolution of the actuation on the 6 joints over a group of teleoperated episodes for the SO-100 manipulator.46Notice how proprioperceptive states are captured jointly with camera frames over the course of the recorded episodes, providing a unified high-frame rate collection of teleoperation data.47Figure the referenced figure shows $(o_t, a_t)$-pairs for the same dataset, with the actions performed by the human expert illustrated just alongside the corresponding observation.48In principle, (expert) trajectories $\tau^\{(i)\}$ can have different lengths since demonstrations might exhibit multi-modal strategies to attain the same goal, resulting in possibly multiple, different behaviors.49 50<ResponsiveImage src={Ch4ObservationActionMapping} alt="Sample observations and action pairs over the course of a given trajectory recorded in /svla_so101_pickplace. Observations, comprising of both proprioperceptive and visual information, are recorded alongside the configuration of a second, leader robot controlled by a human expert, providing complete information for regressing actions given observation." id="fig-fig:ch4-observation-action-mapping" />51 52*Sample observations and action pairs over the course of a given trajectory recorded in /svla_so101_pickplace. Observations, comprising of both proprioperceptive and visual information, are recorded alongside the configuration of a second, leader robot controlled by a human expert, providing complete information for regressing actions given observation.*53 54Behavioral Cloning (BC) [@pomerleauALVINNAutonomousLand1988a] aims at synthetizing synthetic behaviors by learning the mapping from observations to actions, and in its most natural formulation can be effectively tackled as a *supevised* learning problem, consisting of learning the (deterministic) mapping $f: \mathcal\{O\} \mapsto \mathcal\{A\}, \ a_t = f(o_t)$ by solving55$$\begin\{equation\}56 57    \min_\{f\} \mathbb\{E\}_\{(o_t, a_t) \sim p(\bullet)\} \mathcal L(a_t, f(o_t)),58\end\{equation\}$$59for a given risk function $\mathcal L:  \mathcal A \times \mathcal A \mapsto \mathbb\{R\}, \ \mathcal L (a, a^\prime)$.60 61Typically, the expert's joint observation-action distribution $p: \mathcal\{O\} \times \mathcal\{A\} \mapsto [0,1]$ such that $(o,a) \sim p(\bullet)$ is assumed to be unknown, in keeping with a classic Supervised Learning (SL) framework[^1].62However, differently from standard SL's assumptions, the samples collected in $\mathcal D $, correspoding to observations of the underlying $ p$ are *not* i.i.d., as expert demonstrations are collected *sequentially* in trajectories.63In practice, this aspect can be partially mitigated by considering pairs in a non-sequential order---*shuffling* the samples in $\mathcal D $---so that the expected risk under $ p$ can be approximated using MC estimates, although estimates may in general be less accurate.64Another strategy to mitigate the impact of regressing over non-i.i.d. samples relies on the possibility of interleaving BC and data collection [@rossReductionImitationLearning2011], aggregating multiple datasets iteratively.65However, because we only consider the case where a single offline dataset $\mathcal D $ of (expert) trajectories is already available, dataset aggregation falls out of scope.66 67Despite the inherent challenges of learning on non-i.i.d. data, the BC formulation affords several operational advantages in robotics.68First, training happens offline and typically uses expert human demonstration data, hereby severily limiting exploration risks by preventing the robot from performing dangerous actions altogether.69Second, reward design is entirely unnecessary in BC, as demonstrations already reflect human intent and task completion.70This also mitigates the risk of misalignment and specification gaming (*reward hacking*), otherwise inherent in purely reward-based RL [@heessEmergenceLocomotionBehaviours2017].71Third, because expert trajectories encode terminal conditions, success detection and resets are implicit in the dataset.72Finally, BC scales naturally with growing corpora of demonstrations collected across tasks, embodiments, and environments.73However, BC can in principle only learn behaviors that are, at most, as good as the one exhibited by the demonstrator, and thus critically provides no mitigation for the suboptimal decision making that might be enaced by humans.74Still, while problematic in sequential-decision making problems for which expert demonstrations are not generally available---data migth be expensive to collect, or human performance may be inherently suboptimal---many robotics applications benefit from relative cheap pipelines to acquire high-quality trajectories generated by humans, thus justifying BC approaches.75 76<ResponsiveImage src={Ch4IssuesWithBc} alt="Point-wise policies suffer from limitations due to (A) covariate shifts and poor approximation of (B) multimodal demonstrations. (A) Initially small errors may drive the policy out of distribution, incuring in a vicious circle ultimately resulting in failure. (B) Both modes of reaching for a target object in a scene, either left or right-first, are equally as good and thus equally as likely to be present in a dataset of human demonstrations, ultimately resulting in multimodal demonstrations." id="fig-fig:ch4-issues-with-bc" />77 78*Point-wise policies suffer from limitations due to (A) covariate shifts and poor approximation of (B) multimodal demonstrations. (A) Initially small errors may drive the policy out of distribution, incuring in a vicious circle ultimately resulting in failure. (B) Both modes of reaching for a target object in a scene, either left or right-first, are equally as good and thus equally as likely to be present in a dataset of human demonstrations, ultimately resulting in multimodal demonstrations.*79 80While conceptually elegant, point-estimate policies $ f : \mathcal\{O\} \mapsto \mathcal\{A\}$ learned by solving the referenced figure have been observed to suffer from (1) compounding errors [@rossReductionImitationLearning2011] and (2) poor fit to multimodal distributions [@florenceImplicitBehavioralCloning2022, keGraspingChopsticksCombating2020].81Figure the referenced figure illustrates these two key issues related to learning *explicit policies* [@florenceImplicitBehavioralCloning2022].82Besides sequentiality in $\mathcal D$, compounding errors due to *covariate shift* may also prove catastrophic, as even small $\epsilon$-prediction errors $0 < \Vert \mu(o_t) - a_t \Vert \leq \epsilon $ can quickly drive the policy into out-of-distribution states, incuring in less confident generations and thus errors compounding (Figure the referenced figure, left).Moreover, point-estimate policies typically fail to learn *multimodal* targets, which are very common in human demonstrations solving robotics problems, since multiple trajectories can be equally as good towards the accomplishment of a goal (e.g., symmetric grasps, Figure the referenced figure, right).83In particular, unimodal regressors tend to average across modes, yielding indecisive or even unsafe commands [@florenceImplicitBehavioralCloning2022].84To address poor multimodal fitting, @florenceImplicitBehavioralCloning2022 propose learning the generative model $ p(o, a)$ underlying the samples in $\mathcal D $, rather than an explicitly learning a prediction function $ f(o) = a $.85 86## A (Concise) Introduction to Generative Models87 88Generative Models (GMs) aim to learn the stochastic process underlying the very generation of the data collected, and typically do so by fitting a probability distribution that approximates the unknown *data distribution*, $ p$.89In the case of BC, this unknown data distribution $p$ represents the expert's joint distribution over $(o, a)$-pairs.90Thus, given a finite set of $N$ pairs $\mathcal D = \{ (o,a)_i \}_\{i=0\}^N $ used as an imitation learning target (and thus assumed to be i.i.d.), GM seeks to learn a *parametric* distribution $ p_\theta(o,a)$ such that (1) new samples $(o,a) \sim p_\theta(\bullet)$ resemble those stored in $\mathcal D $, and (2) high likelihood is assigned to the observed regions of the unobservable $ p$.91Likelihood-based learning provides a principled training objective to achieve both objectives, and it is thus extensively used in GM [@prince2023understanding].92 93### Variational Auto-Encoders94 95<ResponsiveImage src={Ch4TaskEffectOnPairs} alt="Intuitively, latent variable in a single latent model may contain information regarding the task being performed, which directly results in the likelihood of the same observation-action pair being different for two different tasks. When (A) picking a block the likelihood of a wide gripper's opening should be higher than narrower one, while it should be the opposite when (B) pushing the block." id="fig-fig:ch4-task-effect-on-pairs" />96 97*Intuitively, latent variable in a single latent model may contain information regarding the task being performed, which directly results in the likelihood of the same observation-action pair being different for two different tasks. When (A) picking a block the likelihood of a wide gripper's opening should be higher than narrower one, while it should be the opposite when (B) pushing the block.*98 99A common inductive bias used in GM posits samples $(o,a)$ are influenced from an unobservable latent variable $z \in Z$, resulting in100$$\begin\{equation\}101 102    p (o,a) = \int_\{\text\{supp\}\{Z\}\} p(o,a \vert z) p(z)103\end\{equation\}$$104Intuitively, in the case of observation-action pairs $(o, a)$ for a robotics application, $z $ could be some high level representation of the underlying task being performed by the human demonstrator.105In such case, treating $ p(o,a)$ as a marginalization over $\text\{supp\}\{Z\}$ of the complete joint distribution $p(o,a,z)$ natively captures the effect different tasks have on the likelihood of observation-action pairs.106Figure the referenced figure graphically illustrates this concept in the case of a (A) picking and (B) pushing task, for which, nearing the target object, the likelihood of actions resulting in opening the gripper---the higher $q_6$, the wider the gripper's opening---should intuitively be (A) high or (B) low, depending on the task performed.107While the latent space $Z $ typically has a much richer structure than the set of all actual tasks performed, the referenced figure still provides a solid framework to learn joint distribution conditioned on unobservable yet relevant factors.108Figure the referenced figure represents this framework of latent-variable for a robotics application: the true, $ z$-conditioned generative process on assigns *likelihood* $p((o,a) \vert z)$ to the single $(o,a)$-pair.109Using Bayes' theorem, one can reconstruct the *posterior* distribution on $\text\{supp\}\{Z\}$, $q_\theta(z \vert o,a)$ from the likelihood $p_\theta(o,a \vert z)$, *prior* $p_\theta(z)$ and *evidence* $p_\theta(o,a)$.110VAEs approximate the latent variable model presented in Section eq:BC-latent-variable) using an *approximate posterior* $q_\phi(z \vert o,a)$ while regressing parameters for a parametric likelihood, $p_\theta(o,a \vert z)$ (Figure Section fig:ch4-latent-variable-model).111 112<ResponsiveImage src={Ch4LatentVariableModel} alt="(A) The latent variable model in a robotics application regulates influence between observed ($o,a)$ variables and an unobservable latent variable. (B) VAEs approximate exact latent variable models by means of variational inference." id="fig-fig:ch4-latent-variable-model" />113 114*(A) The latent variable model in a robotics application regulates influence between observed ($o,a)$ variables and an unobservable latent variable. (B) VAEs approximate exact latent variable models by means of variational inference.*115 116Given a dataset $\mathcal D $ consisting of $ N$ i.i.d. observation-action pairs, the log-likelihood of all datapoints under $\theta $ (in Bayesian terms, the *evidence* $ p_\theta(\mathcal D)$) can thus be written as:117 118$$\log p_\theta(\mathcal D) &= \log \sum_\{i=0\}^N p_\theta ((o,a)_i) 119 &= \log \sum_\{i=0\}^N \int_\{\text\{supp\}\{Z\}\} p_\theta((o,a)_i \vert z) p(z) 120 &= \log \sum_\{i=0\}^N \int_\{\text\{supp\}\{Z\}\} \frac\{q_\theta(z \vert (o,a)_i)\}\{q_\theta(z \vert (o,a)_i)\} \cdot p_\theta((o,a)_i \vert z) p(z) 121 &= \log \sum_\{i=0\}^N \mathbb E_\{z \sim p_\theta(\bullet \vert (o,a)_i)\} \left[ \frac\{p(z)\}\{q_\theta(z \vert (o,a)_i)\} \cdot p_\theta((o,a)_i \vert z) \right], $$122 123where we used the referenced figure in the referenced figure, multiplied by $1 = \frac\{q_\theta(z \vert (o,a)_i)\}\{q_\theta(z \vert (o,a)_i)\}$ in the referenced figure, and used the definition of expected value in the referenced figure.124 125In the special case where one assumes distributions to be tractable, $p_\theta (\mathcal D)$ is typically tractable too, and $\max_\theta \log p_\theta(\mathcal D)$ provides a natural target for (point-wise) infering the unknown parameters $\theta $ of the generative model.126Unfortunately, the referenced figure is rarely tractable when the distribution $ p$ is modeled with approximators such as neural networks, especially for high-dimensional, unstructured data.127 128In their seminal work on Variational Auto-Encoders (VAEs), @kingmaAutoEncodingVariationalBayes2022 present two major contributions to learn complex latent-variable GMs on unstructured data, proposing (1) a tractable, variational lower-bound to the referenced figure as an optimization target to jointly learn likelihood and posterior and (2) high-capacity function approximators to model the likelihood $p_\theta(o,a\vert z)$ and (approximate) posterior distribution $q_\phi(z \vert o,a) \approx q_\theta(z \vert o,a)$.129 130In particular, the lower bound on the referenced figure (Evidence LOwer Bound, *ELBO*) can be derived from the referenced figure applying Jensen's inequality---$\log \mathbb\{E\}[\bullet] \geq \mathbb\{E\} [\log (\bullet)]$---yielding:131 132$$\log p_\theta(\mathcal D) &\geq \sum_\{i=0\}^\{N\} \left(133            \mathbb\{E\}_\{z \sim p_\theta(\cdot \vert (o,a)_i)\} \big[ \log p_\theta((o,a)_i \vert z) \big]134            + \mathbb\{E\}_\{z \sim p_\theta(\cdot \vert (o,a)_i)\} \left[ \log \left( \frac\{p(z)\}\{q_\theta(z \vert (o,a)_i)\} \right) \right]135        \right)136 &= \sum_\{i=0\}^\{N\} \left(137            \mathbb\{E\}_\{z \sim p_\theta(\cdot \vert (o,a)_i)\} \big[ \log p_\theta((o,a)_i \vert z) \big]138        - \text\{D\}_\{\text\{KL\}\} \big[ q_\theta(z \vert (o,a)_i) \Vert p(z) \big]139        \right) $$140 141The true, generally intractable posterior $p_\theta (z \vert o,a)$ prevents computing both the expectation and KL divergence terms in the referenced figure, and therefore @kingmaAutoEncodingVariationalBayes2022 propose deriving the ELBO using an *approximate* posterior $q_\phi(z \vert o,a)$, resulting in the final, tractable ELBO objective,142 143$$\text\{ELBO\}_\{\mathcal D\}(\theta, \phi) = \sum_\{i=0\}^\{N\} \left(144            \mathbb\{E\}_\{z \sim q_\phi(\cdot \vert (o,a)_i)\} \big[ \log p_\theta((o,a)_i \vert z) \big]145        - \text\{D\}_\{\text\{KL\}\} \big[ q_\phi(z \vert (o,a)_i) \Vert p(z) \big]146        \right)147        $$148 149From Jensen's inequality, maximizing ELBO results in maximizing the log-likelihood of the data too, thus providing a natural, tractable optimization target.150Indeed, expectations can be estimated using MC estimates from the learned distributions in the referenced figure, while the KL-divergence term can typically be computed in closed-form (1) modeling $q_\phi $ as a Gaussian $ q_\phi(z \vert o,a) = \mathcal N\big(\mu_\phi(o,a), \Sigma_\phi(o,a) \big)$ and (2) imposing a standard Gaussian prior on the latent space, $p(z) = \mathcal N(\mathbf\{0\}, \mathbf\{I\})$.151 152An intuitive explanation of the learning dynamics of VAEs can be given considering the equivalent case of *minimizing the negative ELBO*, which admits a particularly interpretable factorization153 154$$\min_{\theta, \phi} - \text{ELBO}_{\mathcal (o,a) \sim \mathcal D}(\theta, \phi) &= \min_{\theta, \phi}\mathbf{L^{\text{rec}}}(\theta) + \mathbf{L^{\text{reg}}}(\phi) 155 156\mathbf{L^{\text{rec}}}(\theta) &= \mathbb{E}_{z \sim q_\phi(\cdot \vert o,a} \big[ \log p_\theta(o,a \vert z) \big] 157 158\mathbf{L^{\text{reg}}}(\phi) &= \text{D}_{\text{KL}} \big[ q_\phi(z \vert o,a) \Vert p(z) \big] $159 160For any given  (o,a)  pair, the expected value term of the referenced figure is typically computed via MC estimates, resulting in161$-\mathbb{E}_{z \sim q_\phi(\bullet \vert o,a)} \big[ \log p_\theta(o,a \vert z) \big] = \mathbf{L^{\text{rec}}} \approx - \frac{1}{n} \sum_{i=0}^n \log p_\theta(o,a \vert z_i).$$162Assuming $p_\theta(o,a \vert z)$ is parametrized as an isotropic Gaussian distribution with mean $\mu_\theta (z) \in \mathbb R^d$ and variance $\sigma^2$, the log-likelihood thus simplifies to:163$$\log p(o,a \vert z_i) = -\frac\{1\}\{2\sigma^\{2\}\} \big \Vert (o,a)-\mu_\theta(z_i) \big\Vert_2^2 -\frac\{d\}\{2\}\log(2\pi \sigma^\{2\}) \implies \mathbf\{L^\text\{rec\}\} \approx \frac \{1\}\{n\} \sum_\{i=0\}^n \big\Vert (o,a) - \mu_\theta(z_i) \big \Vert^2_2$$164Indeed, it is very common in practice to approximate from the learned likelihood $p_\theta(o,a \vert z)$ as a parametric distribution (e.g. Gaussians) parametrized by some learned vector of coefficients derived from $\mu_\theta (z), \ z \sim p (\bullet)$.165In all such cases, learning a VAE corresponds to optimally *reconstructing* the examples in $\mathcal D $ by minimizing the L2-error---a very common *supervised learning* objective for regression targets---while regularizing the information compression into the latent, as under the common modeling choice $ p(z) = \mathcal N (\mathbf\{0\}, \mathbf\{I\})$ the referenced figure regularizes the posterior limiting the expressivity of $q_\phi(z\vert o,a)$.166 167### Diffusion Models168 169VAEs approximate probability distributions via a *single* latent variable model, assuming the underlying unknown distribution can be factored according to the referenced figure, and solve the variational inference problem of jointly learning the likelihood $p_\theta $ and (approximate) posterior $ q_\phi $ for such model.170In that, the unknown data distribution $ p(o,a)$ is effectively approximated via $\int_Z p(z) p_\theta(o,a \vert z)$, and the underlying generative process reproduced by (1) sampling a latent variable and (2) learning to decode it into a (ideally) high-likelihood sample under the (unknown) $p(o,a)$.171Diffusion Models (DMs) [@hoDenoisingDiffusionProbabilistic2020] are another class of GMs which treat the similar problem of approximating an underlying unknown data distribution---*variational inference*---by *partially* extending VAEs to the case where *multiple* latent variables influence each other and the generative process underlying $o,a$ itself.172In particular, DMs posit the generative process can be decomposed to a series of piece-wise (Markovian) interactions between (latent) variables (Figure Section fig:ch4-many-latents), resulting in173 174$$p(\underbrace\{o,a\}_\{= z_0\}) &= \int_\{\text\{supp\}\{Z_0\}\} \int_\{\text\{supp\}\{Z_1\}\} \hdots \int_\{\text\{supp\}\{Z_T\}\} p(z_0, z_1, \dots z_T) 175 176    p(z_0, z_1, \dots z_T) &= p(z_T) \prod_\{t=0\}^\{T\} p(z_\{t-1\} \vert z_t), $$177 178where we explicitly showed the marginalization over the multiple latents in the referenced figure, and used the law of conditional probability and Markov property in the referenced figure.179 180<ResponsiveImage src={Ch4ManyLatents} alt="HMLV models posit the data generation process is influenced by a stack of Markov-dependent latent variables, with samples from the posterior distribution being progressively higher up in the hierarchy." id="fig-fig:ch4-many-latents" />181 182*HMLV models posit the data generation process is influenced by a stack of Markov-dependent latent variables, with samples from the posterior distribution being progressively higher up in the hierarchy.*183 184Similarily to VAEs, providing an exact interpretation for the latent variables is typically not possible.185Still, one fairly reasonable application-driven intuition is that, by providing a model of the hierarchical, decoupled interaction of latent variables, Hierarchical Markov Latent Variable (HMLV) models attempt to capture the different resolutions at which different conditioning factors intervene, so that in a robotics application for instance, one could naturally distinguish between early-stage trajectory planning ($t \to T $) and fine-grained adjustments ($ t \to 0$).186In that, HMLV models thus provide a framework to perform variational inference via multiple, sequential sampling steps from different higher level distributions instead of approximating the generative process with a single-latent variable model.187DMs are a particular instantiation of HMLV models for which the posterior $q( z_t \vert z_\{t-1\}) = \mathcal N(z_t \sqrt\{1-\beta_t\}, \beta_t \mathbf\{I\})$ for a given $\beta_t \in \mathbb R^+$, thereby iteratively reducing the signal-to-noise ratio as $\beta_t $ increases along the latents hierarchy.188 189Just like VAEs, DMs attemp to learn to reproduce an underlying data distribution $ p (o,a)$ given a collection of i.i.d. samples approximating the model posited to have generated the data in the first place ( Section eq:BC-multi-latent-model-1).190Similarily to VAEs, DMs approximate the process of sampling from the unknown $p(o,a)$ (1) sampling from an easy-to-sample distribution (e.g., Gaussian) and (2) learning to reconstruct high-likelihood samples under the unknown distribution.191However, in stark contrast with VAEs, the easy-to-sample distribution contains *no mutual information* regarding the data distribution $p(o,a)$.192Crucially, as no information from the sample $(o,a)$ (denoted as $z_0 \equiv (o,a)$ for the sake of notation) is assumed to be propagated throughout the chain of latents, the posterior $q(z_t \vert z_\{t-1\})$ assumes a relatively amicable structure in DMs, reducing complexity.193The *true* likelihood $p(z_\{t-1\} \vert z_t)$ is instead typically approximated using the parametrization $p_\theta (z_\{t-1\} \vert z_t)$.194In that, the information contained in the unknwon data distribution is *reconstructed* via a process in which samples from a fixed distribution are turned into (ideally) high-likelihood samples under $p(o,a)$---a process referred to as *denoising*.195 196Under such model, we can express the log-likelihood of an arbitrary sample as[^2]197 198$$\log p_\theta (\underbrace\{o,a\}_\{= z_0\}) =199    &\mathbb\{E\}_\{z_1 \sim q(\bullet \vert z_0)\} \log p_\theta (z_0 \vert z_1) - 200 201    &\mathbb\{E\}_\{z_\{T-1\} \sim q(\bullet \vert z_0)\} \big[ \text\{D\}_\{\text\{KL\}\} (q(z_T \vert z_\{T-1\}) \Vert p(z_T) ) \big] - \notag202 203    &\sum_\{t=1\}^\{T-1\} \mathbb\{E\}_\{(z_\{t-1\}, z_\{t+1\}) \sim q(\bullet \vert z_0)\} \big[ \text\{D\}_\{\text\{KL\}\} (q(z_t \vert z_\{t-1\}) \Vert p_\theta(z_t \vert z_\{t-1\}) ) \big], \notag$$204 205providing an optimization target in the form of $\max_\theta \log p_\theta (\mathcal D)$.206 207In their seminal work on using DMs for variational inference, @hoDenoisingDiffusionProbabilistic2020 introduce major contributions regarding solving $\min_\theta -\log p_\theta(o,a)$.208In particular, @hoDenoisingDiffusionProbabilistic2020 exclusively adopt a fixed *Gaussian* posterior in the form of $q(z_t \vert z_\{t-1\}) = \mathcal\{N\}(\sqrt\{1-\beta_t\}z_\{t-1\}, \beta_t \mathbf I)$.209The choice of adopting Gaussians has profound implications on the generative process modeled.210Indeed, under the (mild) assumption that the variance is sufficiently small $\beta_t \leq \eta, \eta \in \mathbb R^+$, @sohl-dicksteinDeepUnsupervisedLearning2015 proved that the likelihood $p(z_\{t-1\} \vert z_t)$ is Gaussian as well, which allows for the particularly convenient parametrization of the approximate likelihood $p_\theta (x_\{t-1\} \vert x_t) = \mathcal N(\mu_\theta(x_t, t), \Sigma_\theta(x_t,t)), \ t \in [1,T]$, as well as for closed-form tractability of the KL-divergence terms in the referenced figure.211Further, the posterior's structure also enables an analytical description for the distribution of the $t $-th latent variable, $ q(z_t \vert z_0) = \mathcal N (\sqrt\{\bar\{\alpha\}_t\}z_0, (1-\bar\{\alpha\}_t) \mathbf\{I\})$, with $\alpha_t = 1-\beta_t, \ \bar \alpha_t = \prod_\{k=1\}^t \alpha_k $, which conveniently prevents iterative posterior sampling.212 213<ResponsiveImage src={Ch4DiffusionRobotActions} alt="DMs iteratively corrupt samples (left) from an unknown distribution into a quasi-standard Gaussian (center), learning the displacement field (right) that permits to reconstruct samples from the unknown target distribution by iteratively denoising samples of a tractable, easy-to-sample distribution." id="fig-fig:diffusion-robot-actions" />214 215*DMs iteratively corrupt samples (left) from an unknown distribution into a quasi-standard Gaussian (center), learning the displacement field (right) that permits to reconstruct samples from the unknown target distribution by iteratively denoising samples of a tractable, easy-to-sample distribution.*216 217Finally, adopting Gaussian posteriors permits a particularly pleasing interpretation of the dynamics of training DMs [@permenterInterpretingImprovingDiffusion2024].218By using Gaussian posteriors, the hierarchical latent variables effectively lose increasingly more information circa the original (unknown) distribution's sample, $ z_0$, increasingly distributing according to a standard Gaussian and thus containing no information at all (Figure Section fig:diffusion-robot-actions).219Figure the referenced figure illustrates this procedure on a simplified, bidimensional observation-action distribution, where we considered $o=q_2$ and $a=q^h_2$, with $q_2$ representing the robot's *elbow flex* actuation and $q^h_2$ the human teleoperator's robot elbow flex.220 221<ResponsiveImage src={Ch4ActionVsObservationDistribution} alt="A joint action-observation distribution, in the simplified case where the observation is the elbow-flex actuation in a SO-100, and the action is the recorded position for the same joint in the teleoperator arm. The motion recorded being teleoperated, the points distribute along a the diagonal." id="fig-fig:ch4-action-vs-observation-distribution" />222 223*A joint action-observation distribution, in the simplified case where the observation is the elbow-flex actuation in a SO-100, and the action is the recorded position for the same joint in the teleoperator arm. The motion recorded being teleoperated, the points distribute along a the diagonal.*224 225Because the recorded behavior is teleoperated, measurements mostly distribute along the line $a = o + \eta, \eta \sim N(0,1)$, with $\eta$-variability accouting for minor control inconsistencies (Figure Section fig:ch4-action-vs-observation-distribution).226Using Gaussian posteriors---i.e., adding Gaussian noise---effectively simulates a *Brownian motion* for the elements in the distribution's support (in Figure the referenced figure, $\mathcal\{O\} \times \mathcal\{A\}$), whereby information *diffuses away* from the samples, and comparing the diffused samples to the original data points one can derive an estimate of the total displacement induced by diffusion.227Under the only assumption that the likelihood of the diffused samples is low under the original unknown data distribution, then one can effectively approximate the unkwown distribution by learning to *reverse* such displacement.228This key intuition allows to write a simplified training objective:229 230$$231    \mathcal L(\theta) = \mathbb\{E\}_\{t, z_0, \epsilon\} \big[232        \Vert \epsilon - \epsilon_\theta(\sqrt\{\bar \alpha_t\} z_0 + \epsilon \sqrt\{1 - \bar \alpha_t\}, t) \Vert^2 \big], \quad t \sim \mathcal\{U\}(\{1,\dots,T\}), \quad233        z_0 \sim \mathcal\{D\}, \quad234        \epsilon \sim \mathcal\{N\}(\mathbf\{0\},\mathbf\{I\}).$$235 236In this simplified (minimization) objective, the optimization process differs from the referenced figure in that, rather than maxizing $p_\theta$ directly, the parameters $\theta $ of the pairwise likelihood $ p_\theta(z_\{t-1\} \vert z_t)$ are adjusted to *predict the total displacement* $\epsilon $ for a randomly long ($ t \sim \mathcal\{U\}(\\{1,\dots,T\\}$ )) diffusion process starting from a sample of the target distribution.237 238By learning the total displacement from a generally, uninformative corrupted sample obtained diffusing information and a sample from an unknown distribution---significant ($\Vert \epsilon \Vert > 0$) whenever input and target distribution are sufficiently different--- @hoDenoisingDiffusionProbabilistic2020 show that one can approximate the underlying distribution reversing the displacement, *denoising* samples.239Interestingly, under the hypothesis real-world data belongs to a single higher dimensional manifold (Manifold Hypothesis), @permenterInterpretingImprovingDiffusion2024 show that diffusion learns the gradient of a distance function from any off-point manifold (such as perturbed, uniformative samples), and the data manifold itself.240Following this gradient---i.e., denoising a sample from an uninformative distribution---corresponds to projecting back into the manifold, yielding a procedure to sample from unknown distributions by means of Euclidean projection.241Indeed, under the assumption that $p_\theta (z_\{t-1\} \vert z_t)$ is Gaussian, then sampling $z_\{t-1\} \sim p_\theta(\bullet \vert z_\{t\})$ corresponds to computing242 243$$z_\{t-1\} = \frac\{1\}\{\sqrt\{\alpha_t\}\} \left( z_t - \frac\{\beta_t\}\{\sqrt\{1 - \bar\alpha_t\}\} \epsilon_\theta(z_t, t) \right) + \sigma_t \epsilon, \quad \epsilon \sim \mathcal N(\mathbf\{0\}, \mathbf\{I\}), $$244 245thus showing that the lower-level latent variables in a DM can be obtained by iteratively removing noise from the one-step higher order variable, using the noise regressor $\epsilon_\theta(z_t, t)$ learned minimizing the referenced figure.246 247### Flow Matching248 249The posterior parametrization adopted by DMs proved traditionally effective, yet it raised concerns circa its efficiency at inference time, where a possibly large of compute-expensive denoising steps are needed in order to recover a sample from the target distribution.250Flow Matching (FM) [@lipmanFlowMatchingGenerative2023] extends DMs to the general case of arbitrary, parametrized likelihood and posteriors, and in this defines a superseding class of GMs providing a unified framework for learning *continuous transformations* between distributions, encompassing and generalizing DMs.251Instead of a *stochastic, discrete, multi-step* denoising process, FM aims to learn a *deterministic, continuous, differentiable flow* $\psi [0,1] \times Z \mapsto Z $, formalized starting from possibly time-dependent vector field $ v: [0,1] \times Z \mapsto Z $ transporting samples from a simple prior distribution $ p_0$---e.g., a standard Gaussian---to a more complex, potentially unknown data distribution $p_1$ over time.252Note how FM models time $t \in [0,1]$ to be varying continuously while moving away *from* an easy-to-sample distribution $p_0$ *towards* the unknown data-distribution, $p_1$.253This results in a continuous and deterministic trajectory for each sample, which can be more efficient to generate compared to the stochastic paths of DMs.254Formally, FM can be fully characterized by an ordinary differential equation (ODE) relating instantaneous variations of flows with the underlying vector field, and hence providing complete trajectories over the distributions' support when integrating over time,255 256$$\frac\{d\}\{dt\} \psi(z, t) &= v(t, \psi(t, z))257 258    \psi(0, z) &= z$$259 260FM proved very effective in a variety of applications, ranging from image [@esserScalingRectifiedFlow2024] and video generation [@polyakMovieGenCast2025] to robotics control [@black $p_0$ VisionLanguageActionFlow2024].261Most notably, in their introductory work on FM for GM, @lipmanFlowMatchingGenerative2023 show how DMs can be seen as a specific instance of FM where the *conditional* target vector field $u$ approximated by the noise regressor corresponds to262$$\begin\{equation\}263 264    u(t, z\vert z_0) = \frac\{\frac\{d\}\{dt\}\alpha(1-t)\}\{1 - (\alpha(1-t))^2\}(\alpha(1-t)z - z_0), \quad \alpha(t) = e^\{-\frac12 \int_0^t \beta(s) ds\}, \quad \forall z_0 \in \mathcal D265\end\{equation\}$$266Note that the traditional discrete-time noise-scheduler $\{\beta_t\}_\{t=0\}^T$ is now generalized to a continuous map $\beta : [0,1] \mapsto \mathbb R^+$.267Crucially, @lipmanFlowMatchingGenerative2023 prove that by exclusively optimizing the vector field for individual data points $z_0 \in \mathcal D $ individually, one also retrieves the optimal flow to morph the entire support of the initial distribution $ p_0$ into $p_1 \ \text\{s.t.\} \mathcal D \sim p_1$.268 269<ResponsiveImage src={Ch4NormalizingFlows} alt="Probability distributions can be modified applying vector fields resulting in a flow of mass in the support. When acting over time, vector fields can effectively change the distribution's structure." id="fig-fig:ch4-normalizing-flows" />270 271*Probability distributions can be modified applying vector fields resulting in a flow of mass in the support. When acting over time, vector fields can effectively change the distribution's structure.*272 273While the noising schedule of DMs results in a stochastic process that resembles a random walk, FM allows for more general---potentially, deterministic---likelihood and posterior parametrization.274In the FM literature the likelihood and posterior probabilty densities defined along a HMLV model are typically jointly referred to as a *probability path*, where the distributions for successive adjacent transitions in the HMLV model are related by the (normalized) flow between them (Figure Section fig:ch4-normalizing-flows).275The inherent flexibility of FM is one of their key advantages over DMs, as it opens up the possibility of *learning* more efficient paths.276For instance, one can design probability paths inspired by Optimal Transport (OT)---a subdiscipline studying the problem of finding the most efficient way to morph one probability distribution into another.277Probability paths obtained through OT paths tend to be *straighter* than diffusion paths (Figure Section fig:ch4-diffusion-paths-versus-fm), which can lead to faster and more stable training, as well as higher-quality sample generation with fewer steps at inference time.278By avoiding unnecessary backtracking associated with the inherent stochastic nature of both the noising and denoising process in DMs, test-time compute is typically significantly reduced, while retaining comparable results [@lipmanFlowMatchingGenerative2023].279 280<ResponsiveImage src={Ch4DiffusionVsFlowmatching} alt="Compared to diffusion, flow matching distorts distribution along a less randomic pattern, resulting in a clearer interpolation between source and target distribution. The visualization shows an example comparison between these two methods on joint distribution of robot observations and actions over $T=50$ steps." id="fig-fig:ch4-diffusion-paths-versus-fm" />281 282*Compared to diffusion, flow matching distorts distribution along a less randomic pattern, resulting in a clearer interpolation between source and target distribution. The visualization shows an example comparison between these two methods on joint distribution of robot observations and actions over $T=50$ steps.*283 284In practice, FM can be applied to generative modeling by learning a vector field regressor $v_\theta(z, t)$ to approximate a given target vector field $u(t, z)$.285In the particular case of DMs, $u(t, z)$ is defined as in the referenced figure, while in priciple the target vector field can be learned to induce a particular transportation, or fixed according to OT.286Given a sample from the data distribution $z_1 \sim p_1$ and a sample from an easy-to-sample prior $z_0 \sim p_0$, CFM defines a simple path between them using *linear interpolation* between samples $z_t = (1-t)z_0 + t z_1$, resulting in the target vector field $u(t, z_t) = z_1 - z_0$.287Then, a FM model can be trained with the simple regression objective defined as288 289$$290    \mathcal L(\theta) = \mathbb\{E\}_\{t, z_0, z_1\} \big[291        \Vert v_\theta((1-t)z_0 + t z_1, t) - (z_1 - z_0) \Vert^2 \big], \quad t \sim \mathcal\{U\}([0,1]),$$292 293where $z_0 \sim p_0(\bullet)$ and $z_1 \sim p_1(\bullet)$. Note how in the referenced figure---differently from the referenced figure---time is assumed to be varying continuously $t \sim \mathcal U([0,1])$ rather than discretely $t \sim \mathcal U(\\{0,1\\})$, a key property of flow-based models.294The objective in the referenced figure directly regresses the learned vector field onto the simple, straight path connecting a point from the prior and a point from the data, providing a simulation-free training procedure that is both stable and efficient.295At inference time, samples are generated by starting with $z_0 \sim p_0$ and iteratively refined according to $\frac\{dz\}\{dt\} = v_\theta(z_t, t)$ for $t \in [0,1]$---an operation that can be numerically carried out with standard ODE solvers.296 297## Action Chunking with Transformers298 299While GMs prove useful in learning complex, high-dimensional multi-modal distributions, they do not natively address the compouding errors problem characteristic of online, sequential predictions.300In Action Chunking with Transformers (ACT), @zhaoLearningFineGrainedBimanual2023 present an application of VAEs to the problem of learning purely from offline trajectories, introduce a simple, yet effective method to mitigate error compounding, learning high-fidelity autonomous behaviors.301Drawing inspiration from how humans plan to enact atomically sequences of the kind $a_\{t:t+k\}$ instead of single actions $a_t $, @zhaoLearningFineGrainedBimanual2023 propose learning a GM on a dataset of input demonstrations by modeling *action chunks*.302Besides contributions to learning high-performance autonomous behaviors, @zhaoLearningFineGrainedBimanual2023 also introduce hardware contributions in the form of a low-cost bimanual robot setup (ALOHA) capable of performing fine-grained manipulation tasks, such as opening a lid, slotting a battery in its allotment or even prepare tape for application.303 304On the robot learning side of their contributions, @zhaoLearningFineGrainedBimanual2023 adopt transformers as the architectural backbone to learn a *Conditional* VAE [@sohnLearningStructuredOutput2015].305Conditional VAEs are a variation of the more standard VAE formulation introducing a conditioning variable on sampling from the latent prior, allowing the modeling of *one-to-many* relationships between latent and data samples.306Further, in stark contrast with previous work [@florenceImplicitBehavioralCloning2022,jannerPlanningDiffusionFlexible2022], @zhaoLearningFineGrainedBimanual2023 do not learn a full joint $ p_\theta(o,a)$ on observation and actions.307While the *policy* distribution $p_\theta(a \vert o)$ can in principle be entirely described from its joint $p_\theta(o,a)$, it is often the case that the conditional distribution is intractable when using function approximators, as $p_\theta(a \vert o) = \tfrac\{p_\theta(o,a)\}\{\int_\mathcal\{A\} p_\theta(o,a)\}$ and the integral in the denominator is typically intractable.308Instead of modeling the full joint using a vanilla VAE, @zhaoLearningFineGrainedBimanual2023 propose learning a *conditional* VAE [@sohnLearningStructuredOutput2015] modeling the policy distribution directly $p (a \vert o)$.309 310In practice, when learning from demonstrations adopting CVAEs results in a slight modification to the VAE objective in the referenced figure, which is adapted to311 312$$313    \text\{ELBO\}_\{\mathcal D\}(\theta, \phi, \omega) = \sum_\{i=0\}^\{N\} \left(314            \mathbb\{E\}_\{z \sim q_\phi(\cdot \vert o_i, a_i)\} \big[ \log p_\theta(a_i \vert z, o_i) \big]315        - \text\{D\}_\{\text\{KL\}\} \big[ q_\phi(z \vert o_i, a_i) \Vert p_\omega(z \vert o_i) \big]316        \right)$$317 318Notice how in the referenced figure we are now also learning a new set of parameters $\omega$ for the prior distribution in the latent space.319Effectively, this enables conditioning latent-space sampling (and thus reconstruction) during training, and potentially inference, providing useful when learning inherently conditional distributions like policies.320Further, ACT is trained as a $\beta$-CVAE [@higgins2017beta], using a weight of the KL regularization term in the referenced figure as an hyperparameter regulating the information condensed in the latent space, where higher $\beta$ results in a less expressive latent space.321 322In their work, @zhaoLearningFineGrainedBimanual2023 ablated using a GM to learn from human demonstrations compared to a simpler, supervised objective, $\mathcal L_1(a,a^\prime) = \Vert a - a^\prime \Vert_1$.323Interestingly, they found the performance of these two approaches to be comparable when learning from *scripted* demonstrations.324That is, when learning from data collected rolling out a predetermined set of commands $[q^c_0, q^c_1, \dots]$, GM did *not* prove competitive compared to standard supervised learning.325However, when learning from human demonstrations---i.e., from data collected executing commands coming from a human controller $[q^h_0, q^h_1, \dots]$---they found performance (success rate on a downstream task) to be severily (-33.3 The authors also ablate the action chunking paradigm, reporting significant performance gains for performing action chunking (1 To avoid acting openloop, @zhaoLearningFineGrainedBimanual2023 design an inference process consisting in performing inference at every timestep $t $ and then aggregate overlapping chunks using chunks' exponential moving average.326 327<ResponsiveImage src={Ch4Act} alt="Action Chunking with Transformer (ACT), as in [@zhaoLearningFineGrainedBimanual2023]. ACT introduces an action chunking paradigm to cope with high-dimensional multi-modal demonstration data, and a transformer-based CVAE architecture." id="fig-fig:ch4-act" />328 329*Action Chunking with Transformer (ACT), as in [@zhaoLearningFineGrainedBimanual2023]. ACT introduces an action chunking paradigm to cope with high-dimensional multi-modal demonstration data, and a transformer-based CVAE architecture.*330 331In ACT (Figure Section fig:ch4-act), inference for a given observation $ o \in \mathcal O $ could be performed by (1) computing a prior $ p_\omega(z \vert o)$ for the latent and (2) decoding an action chunk from a sampled latent $z \sim p_\omega(\bullet \vert o)$, similarily to how standard VAEs generate samples, with the exception that vanilla VAEs typically pose $p(z\vert o) \equiv p(z) \sim N(\mathbf\{0\}, \mathbf\{I\})$ and thus skip (1).332 333<ResponsiveImage src={Ch4ActEncoder} alt="The CVAE encoder used in ACT. Input action chunks are first embedded and aggregated with positional embeddings, before being processed alongside embedded proprioperceptive information, and a learned `[`CLS] token used to aggregate input level information, and predict the style variable $z $. The encoder is entirely disregarded at inference time." id="fig-fig:ch4-act-encoder" />334 335*The CVAE encoder used in ACT. Input action chunks are first embedded and aggregated with positional embeddings, before being processed alongside embedded proprioperceptive information, and a learned `[`CLS] token used to aggregate input level information, and predict the style variable $ z$. The encoder is entirely disregarded at inference time.*336 337However, the authors claim using a deterministic procedure to derive $z $ may benefit policy evaluation, and thus avoid sampling from the conditional prior at all.338At test time, instead, they simply use $ z = \mathbf\{0\}$, as the conditional prior on $z $ used in training is set to be the unit Gaussian.339At test time, conditioning on the observation $ o$ is instead achieved through explicitly feeding proprioperceptive and visual observations to the decoder, $p_\theta(a \vert z, o)$, while during training $z $ is indeed sampled from the approximate posterior distribution $ p_\phi(z \vert o, a)$, which, however, disregards image observations and exclusively uses proprioperceptive states to form $o $ for efficiency reasons (as the posterior $ q_\phi $ is completely disregarded at test time).340 341<ResponsiveImage src={Ch4ActDecoder} alt="The CVAE decoder used in ACT, comprising of a full encoder-decoder Transformer architecture. Camera observations from all $ n$ camera views are first embedded using pre-trained visual encoders, and then concatenated to the corresponding positional embeddings. Then, alongside embeddings for the proprioperceptive information available and the style variable $z $ retrieved from the CVAE encoder, the Transformer encoder shares the matrices $ K,Q $ with the Transformer decoder, trained to decode fixed position embeddings into action valid chunks." id="fig-fig:ch4-act-decoder" />342 343*The CVAE decoder used in ACT, comprising of a full encoder-decoder Transformer architecture. Camera observations from all $ n$ camera views are first embedded using pre-trained visual encoders, and then concatenated to the corresponding positional embeddings. Then, alongside embeddings for the proprioperceptive information available and the style variable $z $ retrieved from the CVAE encoder, the Transformer encoder shares the matrices $ K,Q $ with the Transformer decoder, trained to decode fixed position embeddings into action valid chunks.*344 345### Code Example: Learning ACT346 347## Diffusion Policy348 349DMs proved very effective in approximating complex highly dimensional distributions, such as distributions over images [@hoDenoisingDiffusionProbabilistic2020] or videos [@polyakMovieGenCast2025], thanks to their inherent capability to deal with multimodal data and training stability.350In Diffusion Policy (DP), @chiDiffusionPolicyVisuomotor2024 present an application of DMs the field of robot learning, leveraging diffusion to model human expert demonstrations in a variety of simulated and real-world tasks.351Similarily to Action Chunking with Transformer [@zhaoLearningFineGrainedBimanual2023], @chiDiffusionPolicyVisuomotor2024 (1) adopt a modified *observation-conditioned target distribution* instead of the full joint $ p(o,a)$ and (2) predict multiple actions into the future instead of a single action.352Besides the intractability of the observations' marginal $p_\theta(o)$ given $p_\theta(o,a)$, DP's rationale for modeling the data distribution via $p_\theta(a \vert o)$ stems from the rather test-time compute intensive nature of diffusion, whereby generating actions *alongside* observations is likely to result in higher complexity and thus a likely larger number of denoising operations, which would prove ultimately pointless considering robotics applications rely on the capability to generate controls rather than reproducing observations.353 354In practice, conditioning on observation data is achieved conditioning the added noise regressor $\epsilon_\theta $ introduced in the referenced figure on a stack of $ T_o$ observations, resulting in the *conditional* simplified diffusion objective355 356$$\mathcal L(\theta) &= \mathbb\{E\}_\{t, a_\{t:t+H_a\}, \epsilon\} \big[357        \Vert \epsilon - \epsilon_\theta(\sqrt\{\bar \alpha_t\} a_\{t:t+T_a\} + \epsilon \sqrt\{1 - \bar \alpha_t\}, t, o_\{t-T_o:t\}) \Vert^2 \big], 358 359        & t \sim \mathcal\{U\}(\{1,\dots,T\}), \quad360        a_\{t:t+T_a\}, o_\{t-T_o:t\} \sim \mathcal\{D\}, \quad361        \epsilon \sim \mathcal\{N\}(\mathbf\{0\},\mathbf\{I\}). \notag$$362 363Notice how in the referenced figure the noise regressor is conditioned both on the latent variable rank $t $ *and* on a stack of previous observations $ o_\{t-T_o:t\}$.364 @chiDiffusionPolicyVisuomotor2024 claim the combination of (1) conditioning on a horizon of previous observations and (2) predicting multiple actions into the future allows DP to *commit to specific modes* in the data at inference time, which proves essential for good performance and avoiding undecisiveness.365 366<ResponsiveImage src={Ch4DiffusionPolicy} alt="The Diffusion Policy archicture, as in [@chiDiffusionPolicyVisuomotor2024]. A stack of $H_o $ previous observations is used as external conditioning to denoise a group of $ H_a $ actions. Conditioning is used at every layer of a U-Net block, and in practice allows to obtain fully-formed action chunks with as little as $ T=10$ denoising steps." id="fig-fig:diffusion-policy-architecture" />367 368*The Diffusion Policy archicture, as in [@chiDiffusionPolicyVisuomotor2024]. A stack of $H_o $ previous observations is used as external conditioning to denoise a group of $ H_a $ actions. Conditioning is used at every layer of a U-Net block, and in practice allows to obtain fully-formed action chunks with as little as $ T=10$ denoising steps.*369 370Figure the referenced figure shows the convolution-based version of the architecture proposed by @chiDiffusionPolicyVisuomotor2024, illustrating inference on a single sample from $\mathcal D $ for simplicity.371An arbitrarily noisy chunk of $ H_a$ actions $\tilde a_\{t:t+H_a\}$ is mapped to a learned high-dimensional space.372Similarily, both image observations and poses are embedded before being aggregated to the action embeddings.373Then, a U-Net [@ronnebergerUNetConvolutionalNetworks2015] is trained to regress the noise added into $\tilde a_\{t:t+H_a\}$, using observation conditioning information at every layer and seeking to optimize the referenced figure.374At inference time, the noise predictor is used to predict the quantity of noise at every $t \in [T, \dots, 0 ]$ and iteratively subtract it from $\tilde a_\{t:t+T_a\}$, reversing the diffusion process simulated in training conditioned on $o_\{t-T_o:t\}$ to predict $a_\{t:t+T_a\}$.375 376Training using 50-150 demos (15-60 minutes of teleoperation data) DP achieves strong performance on a variety of simulated and real-world tasks, including dexterous and deformable manipulation tasks such as sauce pouring and mat unrolling.377Notably, the authors ablated the relevance of using RGB camera streams as input to their policy, and observed how high frame-rate visual observations can be used to attain performance (measured as success rate) comparable to that of state-based policies, typically trained in simulation with priviledged information not directly available in real-world deployments.378As high-frame rate RGB inputs naturally accomodate for dynamic, fast changing environments, @chiDiffusionPolicyVisuomotor2024's conclusion offers significant evidence for learning streamlined control policies directly from pixels.379In their work, @chiDiffusionPolicyVisuomotor2024 also ablate the performance of DP against their baseline against the size of the dataset collected, showing that DP outperforms the considered baseline for every benchmark size considered.380Further, to accelerate inference, @chiDiffusionPolicyVisuomotor2024 employ Denoising Diffusion Implicit Models [@songDenoisingDiffusionImplicit2022], a variant of Denoising Diffusion Probabilistic Models [@hoDenoisingDiffusionProbabilistic2020] (DDPM) adopting a strictly deterministic denoising paradigm (differently from DDPM's natively stochastic one) inducing the same final distribution's as DDPM's, and yet resulting in 10 times less denoising steps at inference time [@chiDiffusionPolicyVisuomotor2024].381Across a range of simulated and real-world tasks, @chiDiffusionPolicyVisuomotor2024 find DPs particularly performant when implementing a transformer-based network as $\epsilon_\theta$, although the authors note the increased sensitivity of transformer networks to hyperparameters and thus explicitly recommend starting out with a simpler, convolution-based architecture for diffusion (Figure Section fig:diffusion-policy-architecture), which are however reported to be biased towards learning low-frequency components [@tancikFourierFeaturesLet2020] and thus may prove more challenging to train with non-smooth action sequences.382 383### Code Example: Learning Diffusion Policies384 385## Optimized Inference386 387Modern visuomotor policies output *action chunks*--sequences $\pi(o_t) = \mathbf\{A\}_t$ with $\mathbf\{A\}_t = \bigl(a_t,a_\{t+1\},\dots,a_\{t+H_a\}\bigr)$ being a sequence of $H_a \gg 1$ low-level commands enqueued in an action queue, originating from an environment observation, $o_t$.388Predicting series of actions instead of single commands proved essential in learning complex, multi-modal behavior [@zhaoLearningFineGrainedBimanual2023,chiDiffusionPolicyVisuomotor2024].389 390Typically, the robot executes the entire action chunk $\mathbf\{A\}_t $, before a new observation $ o_\{t+H_a\}$ is passed to the policy $\pi $ to predict the next chunk.391This results in open-loop inference in between observations captured every $ H_a$ timesteps.392 @zhaoLearningFineGrainedBimanual2023 adopts a different strategy whereby the robot controller interleaves chunk prediction $\mathbf\{A\}_t \gets \pi(o_t)$ and chunk consumption $a_t \gets **PopFront(( \mathbf\{A\}_t $)** , computing a new chunk of actions at every timestep $ t$ and aggregating the predicted chunks on overlapping sections.393While adaptive---every observation at every timestep $o_t$ is processed---such approaches rely on running inference continuously, which can be prohibitive in resource-constrained scenarios, such as edge deployments.394 395A less resource-intensive approach is to entirely exhaust the chunk $\mathbf\{A\}$ before predicting a new chunk of actions, a strategy we refer to as *synchronous* (sync) inference.396Sync inference efficiently allocates computation every $H_a$ timesteps, resulting in a reduced average computational burden at control time.397In contrast, it inherently hinders the responsiveness of robot systems, introducing blind lags due to the robot being *idle* while computing $\mathbf\{A\}$.398 399We directly assess the lack of adaptiveness of robot systems due to acting open-loop, and the presence of lags at runtime by decoupling action chunk prediction $\mathbf\{A\}$ from action execution $a_t \gets **PopFront**(\mathbf\{A\}_t)$, developing an *asynchronous* (async) inference stack (Section alg:async-inference), whereby a $**RobotClient**$ sends an observation $o_t$ to a $**PolicyServer**$, receiving an action chunk $\mathbf\{A\}_t $ once inference is complete (Section fig:ch4-async-inference).400In this, we avoid execution lags by triggering chunk prediction while the control loop is still consuming a previously available queue, aggregating it with the newly incoming queue whenever available.401In turn, async-inference tightens the loop between action prediction and action execution, by increasing the frequency at which observations are processed for chunk prediction.402Crucially, decoupling action prediction from action execution also directly allows to allocate more computational resources on a remote policy server sending actions to the robot client over networks, something which may prove very effective in resource-constrained scenarios such as low-power robots.403 404<ResponsiveImage src={Ch4AsyncInference} alt="**Asynchronous inference**. Illustration of the asynchronous inference stack. Note that the policy can be run on a remote server, possibly with GPUs." id="fig-fig:ch4-async-inference" />405 406***Asynchronous inference**. Illustration of the asynchronous inference stack. Note that the policy can be run on a remote server, possibly with GPUs.*407 408"'409Algorithm:410 411Asynchronous inference control-loop412alg:robotclient413algorithmic[1]414Input: horizon $ T$, chunk size $H_a $, threshold $ g[0,1]$415Init: capture $o_0$; send $o_0$ to PolicyServer;416receive $_0  (o_0)$417$t $ to $ H_a $418$ a_t  \{PopFront\}(_t)$419Execute($a_t $) execute action at step $ t$420$\{|_t|\}\{H_a\} < g $ queue below threshold421capture new observation, $ o_\{t+1\}$422NeedsProcessing $(o_\{t+1\})$ similarity filter, or triggers direct processing423async_handle $\{AsyncInfer\}(o_\{t+1\})$424Trigger new chunk prediction (non blocking)425$\{\}_\{t+1\}  (o_\{t+1\})$ New queue is predicted with the policy426$_\{t+1\}  f(_t,\{\}_\{t+1\})$ aggregate overlaps (if any)427 428NotCompleted(async_handle)429$_\{t+1\}  _t$ No update on queue (inference is not over just yet)430 431algorithmic432alg:async-inference433 434"'435 436#### Implementation details437 438*Async* inference (1) tightens the control loop by capturing observations more often, directly eliminates idle gaps at runtime, and (2) directly allows to run inference on more powerful computational resources than the ones typically available onboard autonomous robotic platforms.439 440Algorithmically, we attain (1) on the **RobotClient**-side by consuming actions from a readily available queue until a threshold condition on the number of remaining actions in the queue ($\vert \mathbf\{A\}_t \vert / H_a < g$) is met. When this condition is triggered, a new observation of the environment is captured and sent to the (possibly remote) **PolicyServer**.441To avoid redundant server calls and erratic behavior at runtime observations are compared in joint-space, and near-duplicates are dropped.442Two observations are considered near-duplicates if their distance in joint-space is under a predetermined threshold, $\epsilon \in \mathbb R_+$.443Importantly, when the queue available to robot client eventually becomes empty, the most recent observation is processed regardless of similarity.444 445Interestingly, the behavior of async inference can be studied analytically. First, let $\ell$ be a random variable modeling the time needed to receive an action chunk $\mathbf\{A\}$ after sending an observation $o $, i.e. the sum of (1) the time to send across the observation $ o$ between the **RobotClient** and **PolicyServer**, $t_\{C \to S\}$ (2) the inference latency on the **PolicyServer**, $\ell_S$ and (3) the time to send $\mathbf\{A\}$ between the **PolicyServer** and **RobotClient**, $t_\{S \to C\}$. Assuming independence, $\mathbb E [\ell] = \mathbb E[t_\{C \to S\}] + \mathbb E[\ell_S] + \mathbb E[t_\{S \to C\}]$ which can be further simplified to $\mathbb E[\ell] \simeq \mathbb E[\ell_S]$, assuming communication time is (1) equal in both directions and (2) negligible with respect to the inference latency. Second, let $\Delta t$ be the environment's control cycle. With a real-world frame-rate of 30 frames per second, $\Delta t=33\text\{ms\}$. Consequently, exhausted queues at runtime--i.e. being idle awaiting for a new chunk--are avoided for $g \geq \frac\{\mathbb E[\ell_S] / \Delta t\}\{H_a\}$. In this, the queue threshold $g$ plays a major role relatively to the availability of actions to the **RobotClient**.446 447the referenced figure illustrates how the size of the action chunk $\lvert \mathbf\{A\}_t \rvert $ evolves over time for three representative values of $ g$, detailing the following key scenarios:448 449- **Sequential limit $(g=0)$.** The client drains the entire chunk before forwarding a new observation to the server. During the round-trip latency needed to compute the next chunk, the queue is empty, leaving the robot *incapable of acting*. This reproduces the behavior of a fully sequential deployment and results in an average of $\mathbb E[\ell_S]$ idle seconds.450 451- **Asynchronous inference $(g \in (0,1))$.** Allowing the client to consume $1-g$ of its available queue $\mathbf\{A\}_\{t-1\}$ before triggering inference for a new action queue $\mathbf\{A\}_\{t\}$, amortizing computation while keeping the queue from emptying. The overlap between successive chunks provides a buffer against modeling errors without the full cost of the $g=1$ regime. The updated queue $\mathbf\{A\}_t$ is obtained aggregating queues on the overlapping timesteps between $\mathbf\{A\}_\{t-1\}$ and the incoming $\tilde\{\mathbf\{A\}\}_\{t\}$.452 453- **Compute-intensive limit $(g=1)$.** As an extreme case, and in keeping with \@zhaoLearningFineGrainedBimanual2023, an observation is sent at *every* timestep. The queue is therefore almost always filled, with only a minor saw-tooth due to$\Delta t/\mathbb E[\ell_s] < 1$. While maximally reactive, this setting incurs one forward pass per control tick and can prove prohibitively expensive on limited hardware. Importantly, because the client is consuming actions while the server computes the next chunk, the available queue never gets filled again.454 455<ResponsiveImage src={Ch4Queues} alt="Action queue size evolution at runtime for various levels of $g $ when (A) not filtering out observation based on joint-space similarity and (B) filtering out near-duplicates observation, measuring their similarity in joint-space." id="fig-fig:ch4-queues" />456 457*Action queue size evolution at runtime for various levels of $ g$ when (A) not filtering out observation based on joint-space similarity and (B) filtering out near-duplicates observation, measuring their similarity in joint-space.*458 459the referenced figure emphasizes the trade-off governed by $g $: small values place result in idle periods, whereas $ g\approx 1$ assumes a highly accurate model and pays a significant compute price. In practice, choosing $g\in(0,1)$ allows to strike a balance between reactivity against resource budgets.460If not for the aforementioned similarity filter, the **RobotClient** would send observations for processing every $(1 - g) H_a \cdot \Delta t$ seconds, receiving a new chunk of actions every $(1 - g) H_a \cdot \Delta t + \mathbb E[\ell_S]$, on average.461The presence of the observation similarity filter dilates this processing time, and serves the scope of avoiding the robot stalling due to the queue being constantly integrated with an incoming, nearly identical, action chunk.462In particular, the referenced figure results in a queue which is filled with incoming actions *unless* near-duplicate observations are filtered out from the processing pipeline. For clarity, the red arrow in the referenced figure highlights a timestep where the observation similarity mechanism is bypassed, forcing a (nearly identical) observation to be processed as the queue results empty.463 464### Code Example: Using Async Inference465 466[^1]: Throughout, we will adopt the terminology and notation for SL introduced in @shalev-shwartzUnderstandingMachineLearning2014467 468[^2]: $o,a = z_0$ for the sake of notation. Steps omitted for brevity. See Section A in @hoDenoisingDiffusionProbabilistic2020 for a complete derivation.