站长平台怎么添加网站,云南网站建设企业推荐,外贸网站推广中山,酒泉建设局造价官网站前言
本文一开始是属于此文《UMI——斯坦福刷盘机器人#xff1a;从手持夹持器到动作预测Diffusion Policy(含代码解读)》的第三部分#xff0c;考虑后Diffusion Policy的重要性很高#xff0c;加之后续还有一系列基于其的改进工作
故独立成本文#xff0c;且写的过程中 …前言
本文一开始是属于此文《UMI——斯坦福刷盘机器人从手持夹持器到动作预测Diffusion Policy(含代码解读)》的第三部分考虑后Diffusion Policy的重要性很高加之后续还有一系列基于其的改进工作
故独立成本文且写的过程中
发现Diffusion Policy的实现代码值得好好剖析下故直接在本文的第二部分 分析Diffusion Policy的代码实现 代码分析会侧重尽可能一目了然、一看就懂如此结合原理和对应的代码实现会理解更深刻把原属于另一篇文章的「Diff-Control、ControlNet详解」已归纳至本文作为第三部分 第一部分 Diffusion Policy
1.1 什么是Diffusion Policy及其优势
1.1.1 什么是扩散策略
所谓扩散策略是指将机器人的视觉运动策略表示为条件去噪扩散过程来生成机器人行为的新方法如下图所示 a)具有不同类型动作表示的显式策略(Explicit policy with different types of action representations)b)隐式策略学习以动作和观察为条件的能量函数并优化能够最小化能量景观的动作Implicit policy learns an energy function conditioned on both action and observation and optimizes for actions that minimize the energy landscapec)通过“条件去噪扩散过程在机器人行动空间上生成行为”即该扩散策略策略不直接输出一个动作而是推断出「基于视觉观察的动作-评分梯度」进行K次去噪迭代instead of directly outputting an action, the policy infers the action-score gradient, conditioned on visual observations, for K denoising iterations
1.1.2 扩散策略的优势表达多模态动作分布、高维输出空间、稳定训练
扩散策略有以下几个关键特性
可以表达多模态动作分布 通过学习动作评分函数的梯度Song和Ermon2019并在此梯度场上执行随机朗之万动力学采样扩散策略可以表达任意可归一化的分布Neal等2011其中包括多模态动作分布By learning the gradient of the action score function Song and Ermon (2019) and performing Stochastic Langevin Dynamics sampling on this gradient field,Diffusion policy can express arbitrary normalizabledistributions Neal et al. (2011), which includes mul-timodal action distribution高维输出空间 正如其令人印象深刻的图像生成结果所示扩散模型在高维输出空间中表现出了极好的可扩展性。这一特性使得策略能够联合推断一系列未来动作而不是单步动作这对于鼓励时间上的动作一致性和避免比较短浅的规划至关重要稳定训练 训练基于能量的策略通常需要负采样来估计难以处理的归一化常数这已知会导致训练不稳定(Du et al., 2020;Florence et al., 2021) 扩散策略通过学习能量函数的梯度绕过了这一要求从而在保持分布表达能力的同时实现了稳定训练
进一步为了充分释放扩散模型在物理机器人上进行视觉运动策略学习的潜力作者团队提出了一套关键的技术贡献包括将闭环动作序列、视觉条件化和时间序列扩散transformer结合起来
闭环动作序列 结合策略预测高维动作序列的能力与递推视界控制来实现稳健执行(We combine the policy’s capability to predict high-dimensional action sequences with receding-horizon control to achieve robust execution) 这种设计允许策略以闭环方式持续重新规划其动作同时保持时间上的动作一致性从而在长视距规划和响应性之间实现平衡视觉条件化 引入了一种视觉条件化扩散策略其中视觉观测被视为条件而不是联合数据分布的一部分。在这种公式化中策略只需提取一次视觉表示而无需考虑去噪迭代从而大大减少了计算量并实现了实时动作推理where the visual observations are treated as conditioning instead of apart of the joint data distribution. In this formulation,the policy extracts the visual representation onceregardless of the denoising iterations,which drastically reduces the computation and enable sreal-time action inference时间序列扩散transformer 提出了一种基于transformer的新型扩散网络该网络最小化了典型CNN模型的过度平滑效应并在需要高频动作变化和速度控制的任务中实现了最先进的性能
1.2 Diffusion for Visuomotor Policy Learning
1.2.1 扩散策略背后的扩散过程回顾
众所周知从高斯噪声中采样的开始DDPM「如果此前不了解什么是DDPM的请详见此文《图像生成发展起源从VAE、VQ-VAE、扩散模型DDPM、DETR到ViT、Swin transformer》的第二部分」执行次去噪迭代以生成一系列噪声逐渐减少的中间动作直到形成所需的无噪声输出 (说白了就是去噪) 该过程遵循下述所示的公式1 其中为通过学习优化参数的噪声估计网络为每次迭代时加入的高斯噪声且上面的公式1也可以理解为一个单一的噪声梯度下降步长定义为如下公式2 其中噪声估计网络有效地预测了梯度场为学习速率 此外公式1中的、和作为与迭代步长相关的函数选择被称为噪声调度可以理解为梯度下降过程中学习速率的调整策略。经证明将设定略小于1能够改善稳定性再之后训练过程首先从数据集中随机抽取未修改的样本开始 对于每个样本随机选择一个去噪迭代然后为迭代采样一个具有适当方差的随机噪声 然后要求噪声估计网络从添加噪声的数据样本中预测噪声如下公式3 最小化公式3所示的损失函数也同时最小化了数据分布和从DDPM 中提取的样本分布之间KL-散度的变分下界「minimizing the loss functionin Eq 3 also minimizes the variational lower bound of the KL-divergence between the data distribution p(x0) and the distribution of samples drawn from the DDPM q(x0) usingEq 1」
1.2.2 改造DDPM使其表达机器人视觉运动策略
虽然DDPM通常用于图像生成但该团队使用DDPM来学习机器人的视觉运动策略。这需要针对DPPM的公式进行两大修改
之前输出的是图像现在需要输出为机器人的动作(changing the output x to represent robot actions)去噪时所依据的去噪条件为输入观测 (making the denoising processes conditioned on input observation Ot)
为了达到以上这两个修改的目的需要做以下措施
一、闭环动作序列预测 一个有效的动作制定应该鼓励在长期规划中的时间一致性和平滑性同时允许对意外观察做出迅速反应 为了实现这一目标应该在重新规划之前采用扩散模型生成的固定持续时间的行动序列预测即如论文原文所说“An effective action formulation should encourage temporal consistencyand smoothness in long-horizon planning while allowingprompt reactions to unexpected observations. To accomplishthis goal, we commit to the action-sequence prediction produced by a diffusion model for a fixed duration beforere planning” 具体来说在时间步骤策略将最新的个观察数据作为输入并预测个动作其中个动作在不重新规划的情况下在机器人上执行 (在此定义中表示观测范围表示动作预测范围而则代表了动作执行范围) 这样做既促进了时间动作的一致性又保持了响应速度 如下图所示 a) 一般情况下该策略在时间步长时将最新的步观测数据作为输入并输出步动作General formulation. At time step t, the policy takes the latest To steps of observation data Ot as input and outputs Ta steps of actions Atb) 在基于CNN的扩散策略中对观测特征应用FiLM(Feature-wise Linear Modulation)来调节每个卷积层通道In the CNN-based Diffusion Policy, FiLM (Feature-wise Linear Modulation) [Film: Visual reasoning with a general conditioning layer] conditioning of the observation feature Ot is applied to every convolution layer, channel-wise. 然后从高斯噪声中提取的减去噪声估计网络的输出并重复次得到去噪动作序列 「这个过程是扩散模型去噪的本质如不了解DDPM请详看此文《图像生成发展起源从VAE、扩散模型DDPM、DETR到ViT、Swin transformer》」Starting from AtK drawn from Gaussian noise, the outputof noise-prediction network εθ is subtracted, repeating K times to get At0, the denoised action sequence. c) 在基于Transformer的扩散策略中观测的嵌入被传递到每个Transformer解码器块的多头交叉注意力层 每个动作嵌入使用所示注意力掩码进行约束使其仅关注自身和之前的动作嵌入(因果注意力类似GPT的解码策略)In the Transformer-based [52]Diffusion Policy, the embedding of observation Ot is passed into a multi-head cross-attention layer of each transformer decoder block. Each action embedding is constrained to only attend to itself and previous action embeddings (causal attention) using the attention mask illustrated. 二、视觉观察条件化 总之他们使用DDPM来近似条件分布而不是Janner等人[20]用于规划的联合分布「We use a DDPM toapproximate the conditional distribution p(At|Ot) instead of the joint distribution p(At,Ot) used in Janner et al. (2022a)for planning. 」 这种形式使得模型能够在不推断未来状态的情况下预测基于观察的动作(This formulation allows the model to predict actionsconditioned on observations without the cost of inferringfuture states)加快了扩散过程并提高生成动作的准确性 为了获取条件分布将下述公式1 修改为如下公式4 将训练损失由下述公式3 修改为如下的公式5
1.3 几个关键设计涉及模型选型等
第一个设计决策是选择的神经网络架构对此作者研究了两种常见的网络架构类型即CNN和Transformer
1.3.1 噪声估计网络架构的选型基于CNN还是基于transformer
1.3.1.1 CNN-based Diffusion Policy
基于CNN的扩散策略中采用Janner等人[21]的1D temporal CNN并做了一些修改如下图所示 首先仅通过特征线性调制FiLM将动作生成过程条件化在观察特征上 并进行去噪迭代来建模条件分布we only model the conditional distribution p(At|Ot)by conditioning the action generation process on observation features Ot with Feature-wise Linear Modulation (FiLM)其次仅预测动作轨迹而非连接的观察动作轨迹we only predict the action trajectory instead of the concatenated observation action trajectory第三利用receding prediction horizon删除了基于修复的目标状态条件反射。然而目标条件反射仍然是可能的与观测所用的FiLM条件反射方法相同we removed inpainting-based goal state conditioning due to incompatibility with our framework utilizing a receding prediction horizon. However, goal conditioning is still possible with the same FiLM conditioning method used for observations
在实践中发现基于CNN的骨干网络在大多数任务上表现良好且无需过多超参数调优。然而当期望的动作序列随着时间快速而急剧变化时(如velocity命令动作空间)它的表现很差可能是由于时间卷积的归纳偏差[temporal convolutions to prefer lowfrequency signals]以更倾向低频信号所致
1.1.1.2 Time-series diffusion transformer
此外为减少CNN模型中过度平滑效应[49]他们还提出了一种基于Transformer架构、借鉴minGPT[42]思想的DDPM来进行动作预测如下图所示 行动和噪声作为transformer解码器块的输入tokens传入扩散迭代的正弦嵌入作为第一个token(Actions with noise At k are passed in as input tokens for the transformer decoder blocks, with the sinusoidal embedding for diffusion iteration k prepended as the first token) 观测通过共享的MLP转换为观测嵌入序列然后作为输入特征传递到transformer解码器堆栈中(The observation Ot is transformed into observation embedding sequence by a shared MLP, which is then passed into the transformer decoder stack as input features) “梯度”由解码器堆栈的每个对应输出token进行预测(The gradient εθ (Ot ,At k , k) is predicted by each corresponding output token of the decoder stack)在基于状态的实验中大多数性能最佳的策略都是通过Transformer骨干实现的特别是当任务复杂度和动作变化率较高时。然而他们发现Transformer对超参数更敏感 当然了Transformer训练的困难[25]并不是Diffusion Policy所独有的未来可以通过改进Transformer训练技术或增加数据规模来解决However, we found the transformer to be more sensitive to hyperparameters. The difficulty of transformer training [25] is not unique to Diffusion Policy and could potentially be resolved in the future with improved transformer training techniques or increased data scale
故一般来说
建议从基于CNN的扩散策略实施开始作为新任务的第一次尝试In general, we recommend starting with the CNN-based diffusion policy implementation as the first attempt at a new task但如果由于任务复杂性或高速率动作变化导致性能低下那么可以使用时间序列扩散Transformer来潜在地提高性能但代价是额外的调优If performance is low due to task complexity or high-rate action changes, then the Time-series Diffusion Transformer formulation can be used to potentially improve performance at the cost of additional tuning
1.3.2 视觉编码器把图像潜在嵌入化并通过扩散策略做端到端的训练
视觉编码器将原始图像序列映射为潜在嵌入并使用扩散策略进行端到端的训练(The visual encoder maps the raw image sequence into a latent embedding Ot and is trained end-to-end with the diffusion policy)
不同的相机视图使用不同的编码器以对每个时间步内的图像独立编码然后连接形成且使用标准的ResNet-18(未进行预训练)作为编码器并进行以下修改:
使用空间softmax池化替代掉全局平均池化以维护空间信息1) Replace the global average pooling with a spatial softmax pooling to maintain spatial information[29]采用GroupNorm代替BatchNorm来实现稳定训练这个修改对于把归一化层与指数移动平均一起使用时(通常应用于DDPM)很有必要2) Replace BatchNorm with GroupNorm [57] for stabletraining. This is important when the normalization layer isused in conjunction with Exponential Moving Average [17](commonly used in DDPMs)
1.4 扩散策略的稳定性与好处
1.4.1 动作序列预测的好处
由于高维输出空间采样困难在大多数策略学习方法中一般不做序列预测。例如IBC将难以有效地采样具有非光滑能量景观的高维动作空间。类似地BC-RNN和BET难以确定动作分布中存在的模式数量(需要GMM或k-means步骤)
相比之下DDPM在不降低模型表现力的前提下在输出维度增加时仍然保持良好扩展性在许多图像生成应用中已得到证明。利用这种能力扩散策略以高维动作序列的形式表示动作它自然地解决了以下问题
时间动作一致性如下图所示为了将T块从底部推入目标策略可以从左或右绕T块走 然而如果序列中的每个动作被预测为独立的多模态分布(如在BC-RNN和BET中所做的那样)。连续动作可能会从不同模式中提取出来并导致两个有效轨迹之间交替出现抖动动作However, suppose each action in the sequence is predicted as independent multimodal distributions (as done in BCRNN and BET). In that case, consecutive actions could be drawn from different modes, resulting in jittery actions that alternate between the two valid trajectories.对于空闲动作的鲁棒性当演示暂停并导致相同位置或接近零速度的连续动作序列时则会发生空闲行为。这在远程操作等任务中很常见 然而单步策略容易过度适应这种暂停行为。例如在实际世界实验中使用BC-RNN和IBC时经常会卡住未删除训练数据集中包含的空闲行为(BC-RNN andIBC often get stuck in real-world experiments when the idleactions are not explicitly removed from training)
// 待更
1.4.2 扩散模型在训练中的稳定
隐式策略使用基于能量的模型(EBM)表示动作分布(An implicit policy represents the action distribution using an Energy-Based Model (EBM))如下公式6所示 其中是一个难以计算的归一化常数(相对于a)
为了训练用于隐式策略的EBM使用了infonce风格的损失函数它相当于公式6的负对数似然 在实践中负采样的不准确性会导致EBMs的训练不稳定[11,48]
而扩散策略和DDPM通过建模公式6中相同动作分布的得分函数[46]回避了的估计问题 因此扩散策略的推理过程(公式4)和训练过程(公式5)都不涉及对的评估从而使扩散策略的训练更加稳定
// 待更 第二部分 Diffusion Policy的编码实现与源码解读
扩散策路(Diffusion Policy)包括网络的创建、数据处理、训练、推理和模型的序列化与反序列
导入依赖代码开始部分导入了所需的库和模块。算法配置到类的映射函数 (algo_config_to_class)这个函数根据算法配置algo_config来决定使用哪个算法类DiffusionPolicyUNetDex 或 DiffusionPolicyUNet以及传递给算法的额外参数。DiffusionPolicyUNetDex 类这是一个算法类继承自PolicyAlgo。它实现了扩散策略的核心功能包括_create_networks创建和配置网络结构 _adjust_brightness调整图像的亮度 process_batch_for_training处理训练数据 train_on_batch在单个数据批次上训练模型 log_info记录训练信息 reset重置算法状态为环境rollout做准备 get_action根据观察和目标获取策略动作 _get_action_trajectory获取动作轨迹 serialize 和 deserialize序列化和反序列化模型参数DiffusionPolicyUNet 类这个类与DiffusionPolicyUNetDex类似但可能有一些不同的实现细节比如在网络结构配置上DiffusionPolicyUNetDex 类在创建网络时会根据是否是单手或双手操作来设置不同的参数如self.arm_split 和 self.action_arm_dim。这是为了适应不同的机械手臂配置。 DiffusionPolicyUNet 类没有这样的区分它直接使用self.ac_dim作为动作维度再比如在动作维度处理上在DiffusionPolicyUNetDex中有一个专门的处理流程来区分手臂和手部的动作维度这在_get_action_trajectory方法中体现得尤为明显其中会分别处理手臂和手部的动作。 DiffusionPolicyUNet则没有这样的区分它直接处理整个动作序列再比如在动作序列处理上DiffusionPolicyUNetDex 类在get_action方法中会从动作队列中弹出动作并且每次只处理一个动作 DiffusionPolicyUNet 类在get_action方法中会将整个动作序列放入动作队列并逐个处理辅助函数replace_submodules 和 replace_bn_with_gn这些函数用于替换网络中的特定模块例如将BatchNorm替换为GroupNorm辅助类SinusoidalPosEmb、Downsample1d、Upsample1d、Conv1dBlock、ConditionalResidualBlock1D 和 ConditionalUnet1D这些类定义了网络的不同组件如卷积块、下采样、上采样和条件残差块 相当于定义了一个基于UNet的网络结构用于处理时间序列数据如动作和观察。这个网络结构包括多个层每层都有条件残差块和可能的下采样或上采样操作扩散调度器代码中使用了DDPMScheduler和DDIMScheduler这些是扩散模型中用于控制噪声添加和去除的调度器指数移动平均EMA代码中还提到了EMA这是一种用于平滑模型权重的技术可以提高训练的稳定性
本第二部分的目录 如下
第二部分 Diffusion Policy的编码实现与源码解读
2.1 DiffusionPolicyUNetDex实现扩散策略的核心功能
2.1.1 _create_networks创建和配置网络结构(包含观察编码、噪声预测、噪声调度)
2.1.2 _adjust_brightness调整图像帧的亮度
2.1.3 process_batch_for_training处理从数据加载器中采样的数据批次
2.1.4 train_on_batch在单个数据批次上训练模型(含添加噪声、预测噪声)
2.1.5 log_info从train_on_batch方法中获取信息并记录到TensorBoard
2.1.6 reset
2.1.7 get_action获取策略动作输出
2.1.8 _get_action_trajectory生成动作轨迹(接受当前观察和目标返回动作序列)
2.1.9 serialize 和 deserialize
2.2 辅助函数替换网络中的特定模块(例如将BatchNorm替换为GroupNorm)
2.2.1 replace_submodules替换符合条件的子模块
2.2.2 replace_bn_with_gn调用上面的replace_submodules执行相关模块的替换
2.3 UNet for Diffusion
2.3.1 SinusoidalPosEmb实现正弦位置嵌入
2.3.2 Downsample1d与Upsample1d下采样与下采样的实现
2.3.3 Conv1dBlock实现一维卷积块
2.3.4 ConditionalResidualBlock1D实现条件残差块
2.3.4 ConditionalUnet1D实现条件一维U-Net 2.1 DiffusionPolicyUNetDex实现扩散策略的核心功能
2.1.1 _create_networks创建和配置网络结构(包含观察编码、噪声预测、噪声调度)
这个方法负责创建和配置网络结构包括观察编码器obs_encoder和噪声预测网络noise_pred_net。它还设置了噪声调度器noise_scheduler和指数移动平均EMA模型。
设置观察组设置不同的观察组并从配置中获取编码器的参数
def _create_networks(self):创建网络并将它们放置在 self.nets 中。 # 为 MIMO_MLP 设置不同的观察组if label 在 self.obs_shapes.keys():del self.obs_shapes[label]observation_group_shapes OrderedDict()observation_group_shapes[obs] OrderedDict(self.obs_shapes)encoder_kwargs ObsUtils.obs_encoder_kwargs_from_config(self.obs_config.encoder) 创建观察编码器创建观察编码器并将所有 BatchNorm 替换为 GroupNorm以配合EMA 使用 obs_encoder ObsNets.ObservationGroupEncoder(observation_group_shapesobservation_group_shapes,encoder_kwargsencoder_kwargs,)# 重要# 将所有 BatchNorm 替换为 GroupNorm 以配合 EMA 使用# 如果忘记这样做性能会大幅下降obs_encoder replace_bn_with_gn(obs_encoder)obs_dim obs_encoder.output_shape()[0] 设置手部噪声和动作维度设置是否使用手部噪声并根据手部类型单手或双手设 置动作维度 self.use_handnoise Trueif self.obs_shapes[robot0_eef_pos][0] 3: # 单手self.arm_split 3self.action_arm_dim 7else: # 双手self.arm_split 6self.action_arm_dim 14 创建网络对象通过ConditionalUnet1D(下文2.3.4节会详述)创建噪声预测网络并将观察编码器obs_encoder和噪声预测网络noise_pred_net放置在 self.nets中 # 创建网络对象noise_pred_net ConditionalUnet1D(input_dimself.ac_dim,global_cond_dimobs_dim * self.algo_config.horizon.observation_horizon)# 最终架构有两部分nets nn.ModuleDict({policy: nn.ModuleDict({obs_encoder: obs_encoder,noise_pred_net: noise_pred_net})})nets nets.float().to(self.device) 设置噪声调度器根据配置设置噪声调度器 如果启用了DDPM则使用DDPMScheduler # 设置噪声调度器noise_scheduler Noneif self.algo_config.ddpm.enabled:noise_scheduler DDPMScheduler(num_train_timestepsself.algo_config.ddpm.num_train_timesteps,beta_scheduleself.algo_config.ddpm.beta_schedule,clip_sampleself.algo_config.ddpm.clip_sample,prediction_typeself.algo_config.ddpm.prediction_type) 如果启用了DDIM则使用DDIMScheduler elif self.algo_config.ddim.enabled:noise_scheduler DDIMScheduler(num_train_timestepsself.algo_config.ddim.num_train_timesteps,beta_scheduleself.algo_config.ddim.beta_schedule,clip_sampleself.algo_config.ddim.clip_sample,set_alpha_to_oneself.algo_config.ddim.set_alpha_to_one,steps_offsetself.algo_config.ddim.steps_offset,prediction_typeself.algo_config.ddim.prediction_type) 否则抛错 else:raise RuntimeError() 设置 EMA根据配置设置 EMA 模型 # 设置 EMAema Noneif self.algo_config.ema.enabled:ema EMAModel(modelnets, powerself.algo_config.ema.power) 设置属性设置类的属性包括网络、噪声调度器、EMA 模型、动作检查标志、观察队 列和动作队列 # 设置属性self.nets netsself.noise_scheduler noise_schedulerself.ema emaself.action_check_done Falseself.obs_queue Noneself.action_queue None
2.1.2 _adjust_brightness调整图像帧的亮度
调整图像帧的亮度。这个方法确保输入张量在CUDA设备上并为每个批次生成一个随机的亮度调整因子然后应用这个因子并限制值的范围在0到1之间
方法定义 def _adjust_brightness(self, tensor):调整一批图像帧的亮度。期望张量的形状为 [batch_size, frame_sequence, channels, height, width]。每批次的亮度调整是相同的但在批次之间有所不同。Args:tensor (torch.Tensor): 输入张量。Returns:torch.Tensor: 亮度调整后的张量。 确保张量在 CUDA 设备上检查输入张量是否在 CUDA 设备上。如果不是则抛出一个错误 # 确保张量在 CUDA 设备上if not tensor.is_cuda:raise ValueError(输入张量不是 CUDA 张量) 生成随机亮度调整因子为每个批次生成一个随机的亮度调整因子。因子的范国在 0.8 到1.2 之间但可以根据需要调整范围 # 为每个批次生成一个随机的亮度调整因子# 因子的范围例如在 0.5 到 1.5 之间但可以调整范围batch_size tensor.size(0)brightness_factor torch.empty(batch_size, 1, 1, 1, 1).uniform_(0.8, 1.2).to(tensor.device) 应用亮度调整将亮度调整因子应用到输入张量上 # 应用亮度调整adjusted_tensor tensor * brightness_factor 剪辑值以确保它们在。到1范围内将调整后的张量值剪辑到 0 到1的范围内以确保亮度调整后的值在有效范围内 # 剪辑值以确保它们在 0 到 1 范围内adjusted_tensor.clamp_(0, 1) 返回调整后的张量返回亮度调整后的张量 return adjusted_tensor
2.1.3 process_batch_for_training处理从数据加载器中采样的数据批次
处理从数据加载器中采样的数据批次以过滤相关信息并准备用于训练的批次。它还检查动作是否被归一化到[-1,1]的范围内。
2.1.4 train_on_batch在单个数据批次上训练模型(含添加噪声、预测噪声)
该函数train_on_batch在单个数据批次上训练模型具体而言逐一执行以下步骤
定义 def train_on_batch(self, batch, epoch, validateFalse):在单个数据批次上进行训练Args:batch (dict): 从数据加载器中采样并由 process_batch_for_training 过滤的包含 torch.Tensors 的字典epoch (int): 纪元编号 - 某些算法需要执行分阶段训练和提前停止validate (bool): 如果为 True则不执行任何学习更新Returns:info (dict): 包含相关输入、输出和损失的字典可能与日志记录相关 获取时间步长和动作维度从算法配置中获取观察、动作和预测的时间步长以及动作维度 To self.algo_config.horizon.observation_horizon # 获取观察时间步长Ta self.algo_config.horizon.action_horizon # 获取动作时间步长Tp self.algo_config.horizon.prediction_horizon # 获取预测时间步长action_dim self.ac_dim # 获取动作维度B batch[actions].shape[0] # 获取批次大小 进入无梯度上下文如果 validate 为 True则不计算梯度 with TorchUtils.maybe_no_grad(no_gradvalidate): # 进入无梯度上下文如果 validate 为 True则不计算梯度info super(DiffusionPolicyUNetDex, self).train_on_batch(batch, epoch, validatevalidate) # 调用父类的 train_on_batch 方法actions batch[actions] # 获取动作张量 添加手部噪声如果使用手部噪声并且存在手部观察则生成手部噪声并添加到手部观察中 if self.use_handnoise and robot0_eef_hand in self.obs_shapes.keys(): # 如果使用手部噪声并且存在手部观察handnoise torch.randn_like(batch[obs][robot0_eef_hand]) * 0.03 # 生成手部噪声batch[obs][robot0_eef_hand] handnoise # 将噪声添加到手部观察中 编码观察将观察和目标编码为特征并将观察特征展平 # 编码观察inputs {obs: batch[obs], # 获取观察goal: batch[goal_obs] # 获取目标观察}for k in self.obs_shapes: # 遍历观察形状# 输入的前两个维度应该是 [B, T]assert inputs[obs][k].ndim - 2 len(self.obs_shapes[k])obs_features TensorUtils.time_distributed(inputs, self.nets[policy][obs_encoder], inputs_as_kwargsTrue) # 编码观察assert obs_features.ndim 3 # [B, T, D]obs_cond obs_features.flatten(start_dim1) # 将观察特征展平 采样噪声并添加到动作中生成噪声并根据扩散迭代时间步长将噪声添加到动作中 # 采样噪声以添加到动作中noise torch.randn(actions.shape, deviceself.device) # 生成噪声# 为每个数据点采样一个扩散迭代timesteps torch.randint(0, self.noise_scheduler.config.num_train_timesteps,(B,), deviceself.device).long() # 生成扩散迭代时间步长# 根据每次扩散迭代的噪声幅度将噪声添加到干净的动作中# 这是前向扩散过程noisy_actions self.noise_scheduler.add_noise(actions, noise, timesteps) # 添加噪声到动作中 预测噪声残差使用噪声预测网络noise_pred_net 预测噪声残差 # 预测噪声残差noise_pred self.nets[policy][noise_pred_net](noisy_actions, timesteps, global_condobs_cond) # 预测噪声残差noise_arm noise[..., :self.action_arm_dim].contiguous() # 获取手臂噪声noise_pred_arm noise_pred[..., :self.action_arm_dim].contiguous() # 获取预测的手臂噪声noise_hand noise[..., self.action_arm_dim:].contiguous() # 获取手部噪声noise_pred_hand noise_pred[..., self.action_arm_dim:].contiguous() # 获取预测的手部噪声 计算 L2 损失计算手臂和手部的L2 损失 # L2 损失loss F.mse_loss(noise_pred_arm, noise_arm) * 0.7 F.mse_loss(noise_pred_hand, noise_hand) * 0.3 # 计算 L2 损失 日志记录记录 L2 损失并将损失从计算图中分离 # 日志记录losses {l2_loss: loss # 记录 L2 损失}info[losses] TensorUtils.detach(losses) # 将损失从计算图中分离 梯度步骤如果不是验证模式则执行梯度步骤并更新 EMA if not validate: # 如果不是验证模式# 梯度步骤policy_grad_norms TorchUtils.backprop_for_loss(netself.nets, # 网络optimself.optimizers[policy], # 优化器lossloss, # 损失)# 更新模型权重的指数移动平均值if self.ema is not None: # 如果启用了 EMAself.ema.step(self.nets) # 更新 EMAstep_info {policy_grad_norms: policy_grad_norms # 记录梯度范数}info.update(step_info) # 更新信息字典 返回信息字典返回包含相关输入、输出和损失的信息字典 return info # 返回信息字典
2.1.5 log_info从train_on_batch方法中获取信息并记录到TensorBoard
处理训练批次的信息以汇总要传递给TensorBoard进行记录的信息。
2.1.6 reset
重置算法状态为环境rollout做准备。这包括设置观察队列和动作队列
2.1.7 get_action获取策略动作输出
get_action这个方法用于获取策路动作输出。它接受当前观察和可选的目标作为输入并返回动作张量
def get_action(self, obs_dict, goal_dictNone):获取策略动作输出。Args:obs_dict (dict): 当前观察 [1, Do]goal_dict (dict): 可选目标Returns:action (torch.Tensor): 动作张量 [1, Da]
具体步骤包括
获取观察和动作的时问步长从算法配置中获取观察和动作的时间步长 (注如上文1.2.2节中所述 表示观测范围表示动作预测范围而则代表了动作执行范围」 # obs_dict: key: [1,D]To self.algo_config.horizon.observation_horizonTa self.algo_config.horizon.action_horizon 检查动作队列是否为空如果动作队列为空则运行推理以生成动作序列并将其放入动作队列中 至于其中的_get_action_trajectory下一节会讲其实现 if len(self.action_queue) 0:# 没有剩余动作运行推理# 将 obs_queue 转换为张量字典在 T 维度上连接# import pdb; pdb.set_trace()# obs_dict_list TensorUtils.list_of_flat_dict_to_dict_of_list(list(self.obs_queue))# obs_dict_tensor dict((k, torch.cat(v, dim0).unsqueeze(0)) for k,v in obs_dict_list.items())# 运行推理# [1,T,Da]action_sequence self._get_action_trajectory(obs_dictobs_dict)# 将动作放入队列self.action_queue.extend(action_sequence[0][3:13]) 执行动作从动作队列中取出一个动作并将其维度扩展为[1Da]然后返回该动作 # 有动作从左到右执行# [Da]action self.action_queue.popleft()# [1,Da]action action.unsqueeze(0)return action
2.1.8 _get_action_trajectory生成动作轨迹(接受当前观察和目标返回动作序列)
这段代码实现了一个名为-get_action_trajectory 的方法用于生成动作轨迹它接受当前观察和可选的目标作为输入并返回动作序列
def _get_action_trajectory(self, obs_dict, goal_dictNone):assert not self.nets.training
具体步骤包括
获取时间步长和动作维度从算法配置中获取观察、动作和预测的时间步长以及动作维度 To self.algo_config.horizon.observation_horizonTa self.algo_config.horizon.action_horizonTp self.algo_config.horizon.prediction_horizonaction_dim self.ac_dim 设置推理时问步长根据配置设置推理时间步长 if self.algo_config.ddpm.enabled is True:num_inference_timesteps self.algo_config.ddpm.num_inference_timestepselif self.algo_config.ddim.enabled is True:num_inference_timesteps self.algo_config.ddim.num_inference_timestepselse:raise ValueError 选择网络选择用于推理的网络。如果启用了 EMA则使用 EMA 的平均模型 # 选择网络nets self.netsif self.ema is not None:nets self.ema.averaged_model 编码观察将观察和目标编码为特征 首先确保输入的前两个维度是 [BT〕然后使用时问分布编码器将输入编码为特征 # 编码观察inputs {obs: obs_dict,goal: goal_dict}for k in self.obs_shapes:# 前两个维度应该是 [B, T]assert inputs[obs][k].ndim - 2 len(self.obs_shapes[k])obs_features TensorUtils.time_distributed(inputs, self.nets[policy][obs_encoder], inputs_as_kwargsTrue)assert obs_features.ndim 3 # [B, T, D]B obs_features.shape[0] 重塑观察将观察特征重塑为 Bobs_horizon * obs_dim的形状 # 将观察重塑为 (B,obs_horizon*obs_dim)obs_cond obs_features.flatten(start_dim1) 从高斯噪声初始化动作从高斯噪声初始化动作 # 从高斯噪声初始化动作noisy_action torch.randn((B, Tp, action_dim), deviceself.device)naction noisy_action 初始化调度器初始化噪声调度器 # 初始化调度器self.noise_scheduler.set_timesteps(num_inference_timesteps) 预测噪声并进行逆扩散步骤在每个时间步长上预测噪声并进行逆扩散步骤以去除噪声 for k in self.noise_scheduler.timesteps:# 预测噪声noise_pred nets[policy][noise_pred_net](samplenaction,timestepk,global_condobs_cond)# 逆扩散步骤去除噪声naction self.noise_scheduler.step(model_outputnoise_pred,timestepk,samplenaction).prev_sample 处理动作并返回使用动作时间步长处理动作并返回动作序列 # 使用 Ta 处理动作start To - 1end start Taaction naction[:, start:end]return action
2.1.9 serialize 和 deserialize
序列化和反序列化模型参数。serialize 方法返回一个包含网络参数的字典而 deserialize 方法从这个字典中加载模型参数
2.2 辅助函数替换网络中的特定模块(例如将BatchNorm替换为GroupNorm)
2.2.1 replace_submodules替换符合条件的子模块
该函数用于替换符合条件的子模块
方法定义 def replace_submodules(root_module: nn.Module, predicate: Callable[[nn.Module], bool], func: Callable[[nn.Module], nn.Module]) - nn.Module:Replace all submodules selected by the predicate withthe output of func.predicate: Return true if the module is to be replaced.func: Return new module to use. 它接受三个参数 。root_module根模块类型为 nn. Module 。predicate 谓词函数接受一个模块作为输入返回一个布尔值表示该模块是 否需要被替换 。func函数接受一个模块作为输入返回一个新的模块用于替换原模块检查根模块是否符合条件如果符合条件则直接返回替换后的模块 if predicate(root_module):return func(root_module) 检查pytorch版本是否大于1.9如果版本过低则抛出导入错误 if parse_version(torch.__version__) parse_version(1.9.0):raise ImportError(This function requires pytorch 1.9.0) 查找符合条件的子模块 bn_list [k.split(.) for k, m in root_module.named_modules(remove_duplicateTrue) if predicate(m)] 替换符合条件的子模块具体而言 首先获取父模块 for *parent, k in bn_list:parent_module root_moduleif len(parent) 0:parent_module root_module.get_submodule(..join(parent)) 其次获取源模块 if isinstance(parent_module, nn.Sequential):src_module parent_module[int(k)]else:src_module getattr(parent_module, k) 接着使用func函数生成目标模块 tgt_module func(src_module) 最后将源模块替换为目标模块 if isinstance(parent_module, nn.Sequential):parent_module[int(k)] tgt_moduleelse:setattr(parent_module, k, tgt_module) 验证所有模块已被替换 # verify that all modules are replacedbn_list [k.split(.) for k, m in root_module.named_modules(remove_duplicateTrue) if predicate(m)]assert len(bn_list) 0 这部分代码再次查找所有符合条件的子模块并确保它们已被替换如果还有未被替换的模块则抛出断言错误 返回根模块 return root_module
2.2.2 replace_bn_with_gn调用上面的replace_submodules执行相关模块的替换
该函数用于将所有的BatchNorm层替换为GroupNorm层
方法定义其接受两个参数 def replace_bn_with_gn(root_module: nn.Module, features_per_group: int16) - nn.Module:将所有 BatchNorm 层替换为 GroupNorm 层。Args:root_module (nn.Module): 根模块features_per_group (int): 每组的特征数默认为 16Returns:nn.Module: 替换后的根模块 root_module根模块类型为 nn.Module。 features_per_group每组的特征数默认为 16调用replace_submodules函数 1.根模块传递根模块 root_module replace_submodules(root_moduleroot_module, 2.谓词函数传递一个匿名函数 Lambda x: isinstanceCx, nn. BatchNorm2d)用于检查模块是否为 nn. BatchNorm2d predicatelambda x: isinstance(x, nn.BatchNorm2d), 3.替换函数传递一个匿名函数 Lambda x:nn. GroupNorm(num_groupsx.num_features // features_per_group,num-channel sx.num_features用于将 BatchNorm 层替换为 GroupNorm层 funclambda x: nn.GroupNorm(num_groupsx.num_features // features_per_group, num_channelsx.num_features))
2.3 UNet for Diffusion
2.3.1 SinusoidalPosEmb实现正弦位置嵌入
SinusoidalPosEmb这个类实现了正弦位置嵌入用于将输入张量转换为正弦和余弦形式的嵌入
class SinusoidalPosEmb(nn.Module):// 初始化方法接收一个参数dim表示嵌入的维度并将其存储在实例变量self.dim中def __init__(self, dim):super().__init__()self.dim dim// 前向传播方法接受一个输入张量xdef forward(self, x):device x.devicehalf_dim self.dim // 2emb math.log(10000) / (half_dim - 1)// 生成嵌入向量emb torch.exp(torch.arange(half_dim, devicedevice) * -emb)// 输入张量和嵌入向量相乘emb x[:, None] * emb[None, :]// 将正弦和余弦嵌入拼接在一起emb torch.cat((emb.sin(), emb.cos()), dim-1)return emb
2.3.2 Downsample1d与Upsample1d下采样与下采样的实现
Downsample1d这个类则实现了一维下采样永远将输入张量的时间维度减半
初始化 def __init__(self, dim):super().__init__()self.conv nn.Conv1d(dim, dim, 3, 2, 1) 初始化方法接受一个参数 dim表示输入和输出的通道数并创建一个一维卷积层 self.conv用于下采样 前向传播 def forward(self, x):return self.conv(x) 前向传播方法接受一个输入张量 x并通过卷积层 self.conv 进行下采样然后返回 下采样后的张量
Upsample1d这个类则实现了一维上采样用于将输入张量的时间维度加倍
初始化 def __init__(self, dim):super().__init__()self.conv nn.ConvTranspose1d(dim, dim, 4, 2, 1) 初始化方法接受一个参数dim表示输入和输出的通道数并创建一个一维转置卷积层 self.conv 用于上采样前向传播 def forward(self, x):return self.conv(x)
2.3.3 Conv1dBlock实现一维卷积块
class Conv1dBlock(nn.Module):Conv1d -- GroupNorm -- Mish一维卷积 -- 组归一化 -- Mish 激活函数def __init__(self, inp_channels, out_channels, kernel_size, n_groups8):super().__init__() # 调用父类的初始化方法# 使用 nn.Sequential 创建一个顺序容器self.block nn.Sequential(nn.Conv1d(inp_channels, out_channels, kernel_size, paddingkernel_size // 2), # 一维卷积层带有适当的填充nn.GroupNorm(n_groups, out_channels), # 组归一化层nn.Mish(), # Mish 激活函数)def forward(self, x):return self.block(x) # 前向传播方法将输入 x 通过顺序容器 self.block 进行处理
可以看到
• 初始化方法接受四个参数 。inp_channels输入通道数 。out_channels输出通道数 。kernel_size卷积核大小 。n_groups组归一化的组数默认为 8然后初始化方法创建一个顺序容器 self.block包含以下层 1.一维卷积层nn. Conv1d输入通道数为 inp_channels输出通道数为out_channels卷积核大小为 kernel-size填充为 kernel-size // 2 2.组归一化层nn.GroupNorm
2.3.4 ConditionalResidualBlock1D实现条件残差块
这个类ConditionalResidualBlock1D用于实现条件一维残差块
类定义 class ConditionalResidualBlock1D(nn.Module): # 定义 ConditionalResidualBlock1D 类继承自 nn.Module 初始化方法 def __init__(self, in_channels, # 输入通道数out_channels, # 输出通道数cond_dim, # 条件维度kernel_size3, # 卷积核大小默认为 3n_groups8): # GroupNorm 的组数默认为 8super().__init__() # 调用父类的初始化方法 定义卷积块 self.blocks nn.ModuleList([ # 定义卷积块列表Conv1dBlock(in_channels, out_channels, kernel_size, n_groupsn_groups), # 第一个卷积块Conv1dBlock(out_channels, out_channels, kernel_size, n_groupsn_groups), # 第二个卷积块]) 定义条件编码器 # FiLM modulation https://arxiv.org/abs/1709.07871# predicts per-channel scale and biascond_channels out_channels * 2 # 条件通道数为输出通道数的两倍self.out_channels out_channels # 设置输出通道数self.cond_encoder nn.Sequential( # 定义条件编码器nn.Mish(), # Mish 激活函数nn.Linear(cond_dim, cond_channels), # 线性层nn.Unflatten(-1, (-1, 1)) # 取消展平) 定义残差卷积 # make sure dimensions compatibleself.residual_conv nn.Conv1d(in_channels, out_channels, 1) \ # 定义残差卷积如果输入通道数不等于输出通道数则进行卷积if in_channels ! out_channels else nn.Identity() # 否则使用 Identity 层 前向传播方法 def forward(self, x, cond): # 定义前向传播方法x : [ batch_size x in_channels x horizon ]cond : [ batch_size x cond_dim]returns:out : [ batch_size x out_channels x horizon ] 通过第一个卷积块 out self.blocks[0](x) # 通过第一个卷积块 编码条件 embed self.cond_encoder(cond) # 编码条件embed embed.reshape( # 重新调整条件编码的形状embed.shape[0], 2, self.out_channels, 1)scale embed[:,0,...] # 获取缩放因子bias embed[:,1,...] # 获取偏置out scale * out bias # 应用 FiLM 调制 通过第二个卷积块 out self.blocks[1](out) # 通过第二个卷积块 添加残差连接 out out self.residual_conv(x) # 添加残差连接return out # 返回输出
2.3.4 ConditionalUnet1D实现条件一维U-Net
ConditionalUnet1D这个类继承自nn.Module用于实现条件一维U-Net
首先其接受多个参数用于配置 U-Net 的结构和参数 def __init__(self, input_dim, # 动作的维度global_cond_dim, # 全局条件的维度通常是 obs_horizon * obs_dimdiffusion_step_embed_dim256, # 扩散迭代 k 的位置编码大小down_dims[256,512,1024], # 每个 U-Net 层的通道大小数组的长度决定了层数kernel_size5, # 卷积核大小n_groups8 # GroupNorm 的组数): 初始化父类和设置维度 super().__init__() # 调用父类的初始化方法all_dims [input_dim] list(down_dims) # 设置所有层的维度start_dim down_dims[0] # 设置起始维度 定义扩散步骤编码器用于对扩散步骤进行编码并计算条件维度 dsed diffusion_step_embed_dim # 扩散步骤编码维度diffusion_step_encoder nn.Sequential( # 定义扩散步骤编码器SinusoidalPosEmb(dsed), # 正弦位置嵌入nn.Linear(dsed, dsed * 4), # 线性层nn.Mish(), # Mish 激活函数nn.Linear(dsed * 4, dsed), # 线性层)cond_dim dsed global_cond_dim # 条件维度 定义中间模块使用条件残差块 in_out list(zip(all_dims[:-1], all_dims[1:])) # 输入输出维度对mid_dim all_dims[-1] # 中间层维度self.mid_modules nn.ModuleList([ # 定义中间模块ConditionalResidualBlock1D(mid_dim, mid_dim, cond_dimcond_dim,kernel_sizekernel_size, n_groupsn_groups),ConditionalResidualBlock1D(mid_dim, mid_dim, cond_dimcond_dim,kernel_sizekernel_size, n_groupsn_groups),]) 定义下采样模块使用条件残差块ConditionalResidualBlock1D和下采样层Downsample1d down_modules nn.ModuleList([]) # 定义下采样模块列表for ind, (dim_in, dim_out) in enumerate(in_out): # 遍历输入输出维度对is_last ind (len(in_out) - 1) # 判断是否为最后一层down_modules.append(nn.ModuleList([ # 添加下采样模块ConditionalResidualBlock1D(dim_in, dim_out, cond_dimcond_dim, kernel_sizekernel_size, n_groupsn_groups),ConditionalResidualBlock1D(dim_out, dim_out, cond_dimcond_dim, kernel_sizekernel_size, n_groupsn_groups),Downsample1d(dim_out) if not is_last else nn.Identity() # 如果不是最后一层则添加下采样层否则添加 Identity 层])) 定义上采样模块使用条件残差块ConditionalResidualBlock1D和上采样层Upsample1d up_modules nn.ModuleList([]) # 定义上采样模块列表for ind, (dim_in, dim_out) in enumerate(reversed(in_out[1:])): # 遍历反转后的输入输出维度对is_last ind (len(in_out) - 1) # 判断是否为最后一层up_modules.append(nn.ModuleList([ # 添加上采样模块ConditionalResidualBlock1D(dim_out*2, dim_in, cond_dimcond_dim,kernel_sizekernel_size, n_groupsn_groups),ConditionalResidualBlock1D(dim_in, dim_in, cond_dimcond_dim,kernel_sizekernel_size, n_groupsn_groups),Upsample1d(dim_in) if not is_last else nn.Identity() # 如果不是最后一层则添加上采样层否则添加 Identity 层])) 定义最终卷积层用于生成输出 final_conv nn.Sequential( # 定义最终卷积层Conv1dBlock(start_dim, start_dim, kernel_sizekernel_size),nn.Conv1d(start_dim, input_dim, 1),) 设置模块属性 self.diffusion_step_encoder diffusion_step_encoder # 设置扩散步骤编码器self.up_modules up_modules # 设置上采样模块self.down_modules down_modules # 设置下采样模块self.final_conv final_conv # 设置最终卷积层print(number of parameters: {:e}.format( # 打印参数数量sum(p.numel() for p in self.parameters()))) 前向传播方法 def forward(self, sample: torch.Tensor, # 输入张量timestep: Union[torch.Tensor, float, int], # 扩散步骤global_condNone): # 全局条件x: (B,T,input_dim)timestep: (B,) 或 int扩散步骤global_cond: (B,global_cond_dim)output: (B,T,input_dim) 接受输入张量 sample sample.moveaxis(-1,-2) # 将输入张量的最后一个维度移动到第二个维度# (B,C,T) 处理扩散步骤 timesteps timestep # 获取扩散步骤if not torch.is_tensor(timesteps): # 如果扩散步骤不是张量timesteps torch.tensor([timesteps], dtypetorch.long, devicesample.device) # 将其转换为张量elif torch.is_tensor(timesteps) and len(timesteps.shape) 0: # 如果扩散步骤是标量张量timesteps timesteps[None].to(sample.device) # 将其扩展为一维张量timesteps timesteps.expand(sample.shape[0]) # 将扩散步骤扩展到批次维度 计算全局特征 global_feature self.diffusion_step_encoder(timesteps) # 计算全局特征if global_cond is not None: # 如果全局条件不为空global_feature torch.cat([ # 将全局特征和全局条件拼接在一起global_feature, global_cond], axis-1) 下采样过程 x sample # 获取输入张量h [] # 定义中间结果列表for idx, (resnet, resnet2, downsample) in enumerate(self.down_modules): # 遍历下采样模块x resnet(x, global_feature) # 通过第一个条件残差块x resnet2(x, global_feature) # 通过第二个条件残差块h.append(x) # 将结果添加到中间结果列表x downsample(x) # 通过下采样层 中间过程 for mid_module in self.mid_modules: # 遍历中间模块x mid_module(x, global_feature) # 通过中间模块 上采样过程 for idx, (resnet, resnet2, upsample) in enumerate(self.up_modules): # 遍历上采样模块x torch.cat((x, h.pop()), dim1) # 将当前结果和中间结果拼接在一起x resnet(x, global_feature) # 通过第一个条件残差块x resnet2(x, global_feature) # 通过第二个条件残差块x upsample(x) # 通过上采样层 最终卷积层 x self.final_conv(x) # 通过最终卷积层 最终返回输出张量 x x.moveaxis(-1,-2) # 将输出张量的第二个维度移动到最后一个维度# (B,T,C)return x # 返回输出
第三部分(选读) Diff-Control改进UMI所用的扩散策略(含ControlNet简介)
3.1 Diff-Control是什么及其提出的背景
3.1.1 背景
自从24年年初斯坦福等一系列机器人横空出世以来模仿学习已经成为训练机器人的重要方法其中基于扩散的策略「4-Diffusion policy: Visuomotor policy learning via action diffusion至于扩散策略的详解请见此文《UMI——斯坦福刷盘机器人从手持夹持器到动作预测Diffusion Policy(含代码解读)》的第三部分」——因其有效建模多模态动作分布的能力而脱颖而出从而提升了性能
然而在实践中动作表示的不一致性问题仍然是一个持续的挑战这种不一致性可能导致机器人轨迹分布与底层环境之间的明显差异从而限制控制策略的有效性[5-Robot learning from human demonstrations with inconsistent contexts]
这种不一致性的主要原因通常源于
人类演示的丰富上下文性质[6-What matters in learning from offline human demonstrations for robot manipulation]分布转移问题[7- A reductionof imitation learning and structured prediction to no-regret online learning]以及高动态环境的波动性其实本质上是无状态的缺乏将记忆和先验知识纳入控制器的机制从而导致动作生成的不一致性即they are fundamentally stateless, lackingprovisions for incorporating memory and prior knowledgeinto the controller, potentially leading to inconsistent actiongeneration
先前的方法如
动作分块「8-Learning fine-grained bimanual manipulation with low-cost hardware即ACT其原理详见此文《ACT的原理解析斯坦福炒虾机器人Moblie Aloha的动作分块算法ACT》」和预测闭环动作序列[4-扩散策略]已被提出以解决这一问题此外Hydra[9]和基于航点的操作[10]修改动作表示以确保一致性
然而这些方法通过改变动作表示来解决问题而不是直接使用原始动作
相反能否通过在扩散策略中加入时间转换来明确地施加时间一致性在深度状态空间模型领域[11- Deep state space models for time series forecasting]–[13-How to train your differentiable filter]有效学习状态转换模型能够识别潜在的动态模式
3.1.2 什么是Diff-Control通过ControlNet将状态信息融入扩散策略中
对此作者团队于2024年7月提出了Diff-Control「其对应的论文为Diff-Control: A Stateful Diffusion-based Policy for Imitation Learning」一种基于扩散的状态策略它生成动作并同时学习动作转换模型其基于[14- Adding conditional control to text-to-image diffusion models]在图像生成中引入的ControlNet框架利用它作为转换模型为基础扩散策略提供时间调控
如下图所示先前的动作序列(蓝色)在生成新的动作序列(红色)时用作条件 Diff-Control的关键目标是学习如何将状态信息融入扩散策略的决策过程中
下图展示了这种行为的一个示例 如上图中部所示一个学习近似余弦函数的策略在时间点给定单一观测值时无状态策略(Diffusion Policy)在生成轨迹的延续方面遇到困难。由于存在歧义扩散策略[4]往往会学习到多种模式相比之下如上图左侧所示Diff-Control通过整合时间条件使其能够在生成轨迹时考虑过去的状态 为此所提出的方法利用了最新的ControlNet架构以确保机器人动作生成中的时间一致性 在CV领域中ControlNet用于稳定扩散模型以在生成图像或视频序列时启用额外的控制输入或附加条件
Diff-Control团队将ControlNet的基本原理从图像生成扩展到动作生成并将其用作状态空间模型在该模型中系统的内部状态、观测(摄像头输入)、和人类语言指令共同影响策略的输出「Our method extends the basic principle of ControlNet from image generation to action generation, and use it as a state-space model in which the internal state of the system affects the output of the policy in conjunction with observations (camerainput) and human language instructions」
如下图所示是Diff-Control在“打开盖子”任务中的实际应用 每个时间窗口内(如红色所示)Diff-Control 生成动作序列 Within each time window (depicted in red), Diff-Control generates action sequences在生成后续动作序列时它利用先前的动作作为额外的控制输入如蓝色所示 When generat-ing subsequent action sequences, it utilizes previous actionsas an additional control input, shown in blue. 这种时间过渡是通过贝叶斯公式实现的有效地弥合了独立策略与状态空间建模之间的差距 This temporaltransition is achieved through Bayesian formulation, effec-tively bridging the gap between standalone policies and statespace modeling
3.2 详解ControlNetDiff-Control的构建基础
为了更好的理解Diff-Control接下来咱们来重点讲下ControlNet
3.2.1 给stable diffusion附加上ControlNet类似LoRA微调
在此文《AI绘画原理解析从CLIP、BLIP到DALLE、DALLE 2、DALLE 3、Stable Diffusion》中我们已经了解到可以通过文本提示类似stable diffusion之类的生图工具生成图像然而有的时候想生成符合预期的图像刚有文本提示还不够可能还会给生图工具一些参考的图像相当于文本可以作为prompt条件、参考图像也可以作为prompt条件
而ControlNet可以通过锁定原base模型stable diffusion的参数并制作SD的可训练副本其中可训练副本通过零卷积层与原始锁定SD模型链接而这个零卷积层的权重可以初始化为0具体如下图所示 其中
一个神经模块将特征图x作为输入并输出另一个特征图y 可以表示为 且x和y通常是二维特征图即其中{h, w, c}分别表示特征图中的高度、宽度和通道数为了在这样的模块中添加一个ControlNet锁定原始模块并克隆一个可训练的副本并使用零卷积层连接被锁定的原模型 该零卷积层表示为是一个1×1的卷积层其权重和偏置均初始化为零顺带提一下为何这里用的零卷积呢而非别的某种卷积呢你还真别说作者团队还真曾尝试过用高斯权重初始化的标准卷积层替换零卷积但性能会出现下降具体情况在2.2.3节最后还会解释 当然具体实践时会用两个零卷积实例参数分别为和从而有至于下图右上角的c是希望添加到网络中的条件向量
3.2.2 用于文本到图像扩散的ControlNet——微调扩散模型
Stable Diffusion 本质上是一个包含编码器(如上面说过的SD的原理详见此文)、中间块和跳跃连接解码器的 U-Net 「73-U-net: Convolutional networks for biomedical image segmentation关于U-net的详解详见此文《图像生成发展起源从VAE、VQ-VAE、扩散模型DDPM、DETR到ViT、Swin transformer》的2.1.1 从扩散模型概念的提出到DDPM(含U-Net网络的简介)」 编码器和解码器各包含 12 个块整个模型包含25 个块包括中间块 在这 25 个块中8 个块是下采样或上采样卷积层而另外 17 个块是主块每个主块包含 4 个 resnet 层和 2 个视觉transformer(即ViT)每个 ViT 包含若干交叉注意力和自注意力机制 例如在上图左侧a图 中“SD 编码器块 A” 包含 4 个 resnet 层和 2 个 ViTs而“×3”表示该块重复三次 文本提示被编码为CLIP文本编码器[66-Learning transferable visual models from natural language supervision] 而扩散时间步通过使用位置编码的时间编码器Time Encoder进行编码Text prompts are encoded using the CLIP text encoder [66], and diffusion timesteps are encodedwith a time encoder using positional encoding最终ControlNet结构应用于U-net的每个编码器层级(相当于ControlNet作用在了U-net的每个编码器的副本上即The ControlNet structure is applied to each encoder levelof the U-net)如上图右侧b图的右上角所示 具体来说使用ControlNet创建了Stable Diffusion的12个编码块和1个中间块的可训练副本。这12个编码块分为4种分辨率resolutions64×64,32×32, 16×16, 8×8每种分辨率重复3次 最终这块的输出(通过13个领卷积层)添被加到U-net的12个跳跃连接和1个中间块上the outputs are added to the 12 skip-connections and 1 middle block of the U-net
另我必须负责任的指出
上图b中ControlNet右上角是复制过来的参数称之为可训练的副本上图b中ControlNet右下角那些零卷积层 它们的参数只是一开始被0初始化后续经过反向传播后零卷积层变为非零并影响输出「反向传播更新 ControlNet 中的可训练副本和零卷积层使零卷积权重通过学习过程逐渐过渡到优化值」 对于这点从ControlNet的GitHub主页也可以证实到即Q: If the weight of a conv layer is zero, the gradient will also be zero, and the network will not learn anything. Why zero convolution works? A: This is wrong. Let us consider a very simpleControlNet/blob/main/docs/faq.md上图b中ControlNet右下角的零卷积层会连接 上图b中ControlNet右上角的可训练副本和上图左侧a中Stable Diffusion被冻结的u-net
这种方法加快了训练速度并节省了GPU内存。根据在单个NVIDIA A100PCIE 40GB上的测试使用ControlNet优化Stable Diffusion仅需大约23%的额外GPU内存和34%的计算资源 关于上面的64×64再具体说明下Stable Diffusion 使用类似于 VQ-GAN [19] 的预处理方法将 512×512像素空间图像转换为较小的 64×64 潜在图像 为了将 ControlNet 添加到 Stable Diffusion 中首先将每个输入条件图像例如边缘、姿势、深度等从 512×512 的输入大小转换为 64×64的特征空间向量以匹配 Stable Diffusion 的大小特别地作者团队使用一个小型网络由四个卷积层组成使用 4×4 的卷积核和 2×2 的步幅通过 ReLU 激活分别使用 16、32、64、128通道用高斯权重初始化并与完整模型一起训练将图像空间条件编码为特征空间条件向量 此外条件向量被传入ControlNet 3.2.3 ControlNet的训练、推理、消融实验
首先对于ControlNet的训练
给定一个输入图像图像扩散算法逐步向图像添加噪声并生成一个噪声图像其中表示添加噪声的次数
给定一组条件包括时间步、文本提示以及特定任务条件图像扩散算法学习一个网络来预测添加到噪声图像的噪声 整个扩散模型的总体学习目标是这个学习目标直接用于通过ControlNet微调扩散模型
在训练过程中随机将50%的文本提示替换为空字符串。这种方法增强了ControlNet直接识别输入条件图像如边缘、姿态、深度等语义的能力以替代提示
在训练过程中由于零卷积不会向网络添加噪声模型应始终能够预测高质量的图像。且观察到模型并不是逐渐学习控制条件而是突然成功地遵循输入条件图像通常在不到1万次优化步骤内完成如下图所示一般称之为“突然收敛现象” 其次对于ControlNet的推理
可以通过多种方式进一步控制ControlNet的附加条件如何影响去噪扩散过程
比如无分类器引导的分辨率加权
Stable Diffusion依赖于一种称为无分类器引导CFG的技术来生成高质量的图像 CFG的公式为其中、、、分别是模型的最终输出、无条件输出、有条件输出和用户指定的权重当通过ControlNet添加条件图像时可以将其添加到和或者仅添加到 在具有挑战性的情况下例如没有提示时将其添加到和会完全移除CFG引导下图b仅使用会使引导非常强下图c 作者团队的解决方案是首先将条件图像添加到然后根据每个块的分辨率将权重 乘以 Stable Diffusion 和 ControlNet 之间的每个连接公式为其中是第个块的大小例如 通过降低 CFG 引导强度可以得到上图d 中所示的结果可称之为 CFG 分辨率加权——组合多个 ControlNet为了将多个条件图像例如Canny 边缘和姿态应用于单个 Stable Diffusion 实例还可以直接将相应 ControlNet的输出添加到 Stable Diffusion 模型中如下图所示 这种组合不需要额外的加权或线性插值
最后在消融实验方面(证明两个零卷积结构的必要性)
研究了ControlNets的替代结构
用高斯权重初始化的标准卷积层替换零卷积——如下图b所示用一个单一的卷积层替换每个模块的可训练副本称之为ControlNet-lite——如下图c所示 且提出了4种提示设置以测试现实世界用户可能的行为
无提示不充分的提示这些提示未能完全涵盖条件图像中的对象例如本文的默认提示“高质量、详细且专业的图像”条件图像语义的冲突提示完美提示描述必要的内容语义例如“漂亮的房子”
可以很明显的看到
如上图a所示4种设置中不论是哪一种设置ControlNet的表现都是最好的当零卷积被替换时(比如被替换成高斯权重初始化的标准卷积层)ControlNet的性能下降到与ControlNet-lite大致相同这表明可训练副本的预训练骨干在微调过程中被破坏——上图b轻量级的ControlNet-lite——上图c不足以解释条件图像并在条件不足和无提示的情况下失败
3.3 Diff-Control的技术架构
3.3.1 从扩散模型到递归贝叶斯公式
回顾一下扩散模型的背景知识
扩散模型通过迭代的将高斯噪声映射到目标分布且可以选择性的基于上下文信息进行条件生成
比如给定初始点扩散模型预测输出序列为其中每个后续输出都是前一个输出的去噪版本(说白了整个过程就是去噪的过程)故就是扩散过程的输出然后使用去噪扩散概率模型DDPM作为骨干网络。在训练过程中噪声输入可以通过不同的噪声水平生成其中是方差调度是随机噪声接下来可以训练一个神经网络来预测添加到输入中的噪声通过最小化以下公式 其中表示观察和动作对是去噪时间步 在采样步骤中迭代运行去噪过程 整个过程相当于不断减掉每一步添加的噪声
接下来目标是学习一个以条件和观察为输入的策略
在这个背景下定义为包含机器人末端执行器姿态的轨迹。与之前的方法[4],[32]一致现在的目标也是将多个条件作为输入。然而正如在上文所提到的尽管之前的工作[4], [8], [9]已经努力探索稳健的动作生成但它们并没有考虑到a的状态性
故下面通过在动作空间中引入转移从贝叶斯的角度解决动作一致性问题(We address the action consistency from a Bayesian perspective by introducing transition in action spaces)
首先有如下方程令应用马尔可夫性质即假设下一个生成的轨迹仅依赖于当前的轨迹得到 其中是归一化因子是观测模型是转移模型这个转移模型描述了系统动力学演化的规律而观测模型则确定了系统内部状态与观测到的噪声测量值之间的关系
3.3.2 Diff-Control Policy结合贝叶斯公式与扩散模型
接下来展示如何将贝叶斯公式和扩散模型结合在一起使得一个策略可以生成有状态的动作序列从而促进一致的机器人行为
他们提出了如下图所示的Diff-Control策略其参数为 其中代表执行时间范围表示以自然人类指令形式出现的语言条件表示由场景的RGB相机捕获的一系列图像
策略生成一个轨迹窗口其中指的是窗口大小或预测时间范围
贝叶斯公式中的Diff-Control策略包含两个关键模块
转换模块接收先前的动作并生成潜在的嵌入以供基础策略后续使用作为观测模型基础策略结合了与相关的时间信息并生成一个新的动作
这种双模块结构使得Diff-Control策略能够巧妙地捕捉时间动态并有助于准确且一致地生成后续动作
3.3.3 基础策略与转换模型
对于基础策略如下图左侧所示 首先按照第2.2.1节中的步骤训练一个基于扩散的策略[4]作为基础策略然后采用[15-Planning with diffusion for flexible behavior synthesis]中的一维时间卷积网络并构建U-net骨干网络策略可以自主执行并生成动作而无需依赖任何时间信息
对于转换模型如下图右侧所示作者团队将ControlNet纳入为转换模块「Diff-Control Policy通过使用锁定的 U-net 扩散策略架构来实现。该策略复制了编码器和中间模块并引入了零卷积层即The Diff-Control Policy is implemented through the utilization of a locked U-net diffusion policy architecture. It replicates the encoder and middle blocks and incorporates zero convolution layers」 这种利用有效地扩展了策略网络的能力使其包含时间条件(This utilization extends the capability of the policy networkto include temporal conditioning effectively)
为实现这一目标利用先前生成的动作序列作为ControlNet的提示输入To achieve this,we utilize the previously generated action sequences as the prompt input to ControlNet. 通过这样做基础策略可以了解先前的动作且通过创建一个可训练的编码器副本来实现ControlNet然后冻结基础策略By doing so, the base policy¯πψψψ becomes informed about the previous actions a[Wt−h].We implement ControlNet by creating a trainable replica of the ¯πψψψ encoders and then freeze the base policy ¯πψψψ. 且可训练的副本通过零卷积层[33]连接到固定模型The trainable replica is connected to the fixed model with zero convolutional layers [33]最终ControlNet可以将作为条件向量并重用训练好的基础策略来构建下一个动作序列ControlNet can then take a[Wt−h] as the conditioning vector and reuses the trained base policy¯πψψψ to construct the next action sequence a[Wt]
3.3.4 策略的训练与相关模型的选择
在策略的训练上其中
对于基础策略的训练 首先将观察和语言条件编码到相同的嵌入维度 然后我们利用方程中定义的学习目标来训练基础策略在微调ControlNet时也使用相同的学习目标 整个扩散模型的总体学习目标是是对应的由参数化的神经网络且基础策略和Diff-Control策略都是端到端训练的
在模型的选择上他们采用了基于CNN的U-net架构。这种基于CNN的骨干网络选择已被证明适用于多种任务无论是在模拟环境还是现实场景中
在整个的模型架构中如之前的下图所示 编码器和解码器块使用了具有不同核大小的1D卷积层
为了实现ControlNet的零卷积层采用了权重初始化为零的1D1×1卷积层 这种方法确保在训练的初始阶段任何潜在的有害噪声不会影响可训练神经网络层的隐藏状态[14]且在所有实验中使用窗口大小W24作为默认预测范围执行范围h在和之间为12步故在实验中默认只执行12步h12以与[4-Diffusion policy: Visuomotor policy learning via action diffusion]保持一致 因为根据[4]执行范围过大或过小都可能导致性能下降
此外在实施语言条件任务时他们还观察到基于扩散的策略在利用复杂的CLIP语言特征作为条件学习多样化动作时达到了某种能力上限
为了解决这个问题更实用的方法是引入融合层并增加视觉和语言表示的嵌入大小而不是直接将它们拼接在一起「To address this,a more practical approach is to incorporate a fuse layerand increasing embedding size for visual and language representations, instead of concatenating them directly」这一修改可以提升策略在语言条件任务中的整体表现
3.4 实验与评估围绕4个任务与4种基线方法PK
3.4.1 对4个任务的一般性建模
通过将Diff-Control策略与四种基线方法在五个不同的机器人任务上进行比较包含的任务包括a语言条件下的厨房任务b厨房场景中的开盖任务作为高精度任务c动态场景中的捞鸭任务d作为周期性任务的打鼓任务
UR5机器人手臂的动作表示为其中每个动作表示为其中包括末端执行器在笛卡尔坐标系中的位置、方向和夹爪的关节角度且对于所有任务输入模态包括两种模态和 第一种模态对应于RGB图像 第二种模态指的是从自然语言序列中提取的语言嵌入。这个嵌入作为机器人理解和决策过程的语言输入
下表 按主观难度升序排列任务提供了任务特征的摘要如干扰物数量Dis、专家演示数量Dem、不同动作数量Act以及是否需要高精度HiPrec 以下是具体的4个任务
语言条件厨房任务该任务旨在模拟厨房场景中的多个任务 [35], [36] 机器人工作区由一个缩小的现实模型厨房组成如下图所示 厨房环境包含各种物体包括锅、平底锅、碗和类似毛绒蔬菜的干扰物。在数据收集过程中干扰物会被随机放置。训练有素的专家负责远程操作机器人在厨房环境中执行两个特定动作。这些动作包括取回一个西红柿并根据给定的语言指令将其放置在炉灶上的锅中A或水槽中B高精度开盖任务包括抬起盖子并随后将其放置在附近的碗上这需要精确的控制如下图所示盖子的把手相对较小且盖子的表面是反光的 为了收集此任务的数据他们获得了50次专家示范。每次示范中引入了5个或更多干扰物体的随机放置以及锅的位置和盖子的旋转的轻微变化鸭子舀取受[37]的启发他们为机器人配备了一个勺子机器人的目标是将鸭子从水中舀出来如下图右下角所示这一任务由于勺子进入水中引起的扰动而具有挑战性 水流影响橡皮鸭的位置要求机器人执行精确而谨慎的动作以成功捕捉鸭子鼓点该任务专门为机器人学习周期性动作而设计由于需要独特的动作表示这是一项具有挑战性的任务[38]如下图所示 该任务的难点在于机器人必须准确地计数鼓点的次数并确定何时停止敲鼓。通过远程操作机器人在每次示范中敲击鼓三次共获得了150次专家示范
3.4.2 与4种基线在性能上的PK
为了评估diff-control的泛化能力、整体性能及其优势他们提出的基线如下
Image-BC这个基线采用图像到动作的代理框架类似于BC-Z[2]它基于ResNet-18骨干网络并使用FiLM [39]通过CLIP语言特征进行条件处理ModAttn[32]该方法采用transformer风格的神经网络并使用模块化结构通过神经注意力来处理任务的各个子方面它需要人类专家正确识别每个任务的组件和子任务BC-Z LSTM这个基线代表了一种受BC-Z架构启发的有状态策略。通过使用MLP和LSTM层将先前的动作和语言条件融合实现了先前输入的整合扩散策略[4]这个基线是一个标准的扩散策略
所有基线模型均使用相同的专家演示数据进行再现和训练总共训练3,000个周期
在整个训练过程中每300个周期保存一次检查点。在他们的分析中他们报告了这些保存的检查点中每种基线方法所获得的最佳结果每个实验均在单个NVIDIA Quadro RTX 8000 GPU上以64的批量大小进行持续约24小时对于所有任务使用Adamw优化器学习率为1e-4
具体PK结果详见原论文
