1. 从LeRobot到diffusion module这篇笔记到底在讲什么LeRobot这个框架最近在低成本机械臂圈子里确实火尤其是搭配SO-ARM100这种几百块钱就能搞定的桌面机械臂让不少原本觉得机器人操作高不可攀的人也能跑起模仿学习。我上一篇文章写了怎么把环境搭起来、怎么采集数据这篇笔记直接把镜头对准LeRobot里最核心的动作生成模块——diffusion module也就是扩散策略。先说清楚这篇文章适合谁看。如果你已经在Ubuntu上按过LeRobot知道hdf5和mcap这两种数据集格式的基本区别也跑通过一次简单的训练和推理流程但一直对diffusion模块内部是怎么工作的、为什么它能比普通行为克隆模型生成更平滑的动作感到困惑这篇文章就是写给你看的。如果你刚接触LeRobot也没关系我会把diffusion相关的基础概念用比较直白的话讲清楚但你最好先跑通过官方quickstart不然一些细节会显得比较抽象。在上一篇文章里我提到LeRobot把所有演示数据统一成了两种格式老牌的hdf5和新主推的mcap。diffusion模块的训练数据就来自这些文件但真正用到模型输入时数据要经过观察窗口切分、归一化、转换成tensor等一系列处理。这些处理逻辑藏得比较深很多人第一次打开LeRobot源代码时会被文件跳转搞晕。我这篇笔记就顺着diffusion模块这条线从模型结构、训练流程、推理流程、常见坑这四个维度逐一拆开讲。2. 扩散策略是怎么在机械臂上工作的2.1 扩散模型不是拿来生成图片是拿来生成动作序列如果你接触过AI绘画应该听过扩散模型的大名——先给一张纯噪声图片然后一步步去噪最终还原成一张清晰图片。扩散策略借用了同一个数学框架但去噪的对象从图片换成了机械臂的动作序列。在LeRobot的SO-ARM100场景里一条演示数据是由多个时间步组成的每个时间步包含关节角度、末端位置之类的状态信息。扩散策略的训练目标可以这样理解模型学习一个动作序列的真实分布推理时从一个随机噪声动作序列出发通过多轮去噪迭代逐步逼近一个符合当前观察状态的动作序列。这个概念听起来复杂实际操作上你不需要手写任何去噪算法LeRobot已经把最核心的UNet和DDPM调度器都封装好了。你要做的事情是搞懂几个关键参数知道怎么调才能让你的机械臂动作不抖、不飘、跟演示数据吻合度高。我在跑SO-ARM100的时候发现很多人一上来就纠结“为什么我训练完了机械臂的动作跟抽风一样”“为什么动作很迟钝”其实根子都在diffusion模块的输入输出配置上。先把这个模块的输入输出打通再去调网络结构才会有效果。2.2 观察窗口、预测窗口和动作执行窗口的区别这是diffusion模块里最容易混乱的三个概念我单独拿出来讲。观察窗口obs_horizon每次模型做决策时往前看多少帧的历史观察数据。比如设为2就表示模型看当前帧和上一帧的关节状态、末端位置等。预测窗口pred_horizon模型一次性预测出未来多少步的动作。比如设为16就表示模型输出未来16个控制步的动作序列。动作执行窗口action_horizon既然模型一次预测了16步是不是16步全都要执行不是。执行窗口告诉控制器“你只执行前几步然后重新用最新的观察去预测”。一般是8步也可以设成和pred_horizon一样相当于预测一段执行一段再预测。这三个参数在LeRobot的配置文件里通常写作obs_horizon、pred_horizon、action_horizon。它们在diffusion policy里的作用可以用开车比喻观察窗口是你看后视镜和前方路况的时长预测窗口是你规划接下来几秒的路线执行窗口是你真正踩下油门刹车的持续时间。三者配合得好车就开得稳配合不好要么反应太慢要么动作毛糙。SO-ARM100这种六关节机械臂关节空间维度是6如果pred_horizon设成16那么模型每次要做的事情就是从一个随机的(16, 6)动作序列里逐步去噪最后得到一个(16, 6)的平滑动作序列。这个维度变化直接决定了UNet的输入通道设计后面会讲到。2.3 为什么扩散策略比普通行为克隆更适合机械臂普通的模仿学习做法是输入当前观察输出当前动作本质上是学习一个从观测量到动作量的映射函数。这种做法的局限性在于它假设每个观察点对应的动作是唯一的。但演示数据里经常存在“同一个位置不同轨迹”的情况比如人在两次演示中拿同一杯水时手的路径略有不同普通行为克隆会强行平均这些轨迹导致动作被“磨平”表现就是犹豫、迟钝、不到位。扩散策略学的是动作序列的分布而不是简单映射它天然支持多峰输出。你可以这样理解普通行为克隆是“老师怎么教我就怎么学但老师教了两遍不一样我就傻了”扩散策略是“老师教了两遍不一样我学会了这两条路都能走看情况选哪条”。所以在少量演示数据下扩散策略的鲁棒性明显更好。另外一点很关键扩散策略输出的是动作序列不是单步动作。相当于它做了时间维度的规划动作自然更平滑不会出现高频抖动。这两条结合起来使得LeRobot官方在SO-ARM100的很多任务上默认推荐的就是diffusion策略。3. LeRobot扩散模块的代码结构与关键类拆解3.1 DiffusionPolicy、ConditionalUnet1D、DDIMScheduler三者各自干什么打开LeRobot的源码你会看到lerobot/policies/diffusion.py这个文件里面定义了一个DiffusionPolicy类。这个类是整个扩散策略的门面训练和推理都通过它来调用。它内部主要由三块组成一个ConditionalUnet1D网络、一个噪声调度器默认是DDIM、以及一个用于图像输入的特征提取器如果是视觉策略。ConditionalUnet1D就是那个不断去噪的神经网络。它的输入有两个部分带噪声的动作序列本身以及当前观察状态对应的条件向量。你可以把条件向量理解成“给模型的提示词”告诉它“我现在处于什么状态你要生成符合这个状态的动作”。LeRobot在训练时会对观察状态做归一化然后通过几个全连接层把归一化后的状态映射成条件向量。DDIMScheduler是负责控制噪声添加和去除过程的调度器。训练时它负责往干净动作序列上逐步加噪声推理时它负责从纯噪声序列里逐步去噪。LeRobot默认的扩散步数在训练和推理时可以不同比如训练时设成100步推理时为了速度可以只跑10步DDIM调度器支持这种跳步去噪这也是它比传统DDPM快很多的原因。源码读到这里你会发现一个比较有意思的点在LeRobot的配置里每个策略都有对应的training和deployment配置块。diffusion模块在训练阶段和部署阶段的参数是分开控制的比如训练时的扩散步数、推理时的扩散步数、随机种子等。这个设计让切换实验非常方便不用改核心模型代码改配置就行。3.2 先搞懂归一化再谈训练不然后面全白搭LeRobot在训练diffusion模块前会对所有特征做归一化。这个细节很容易被忽略但它对训练效果的影响非常大。官方在数据集代里会计算每个特征的均值、标准差然后保存下来。训练时网络看到的是归一化后的数据推理时模型输出的是归一化后的动作值要还原成真实关节角度必须经过反归一化。很多人训练完了发现机械臂乱动原因之一就是忘记把模型的输出反归一化回真实值。LeRobot的DiffusionPolicy内部已经封装好了这一步select_action方法会返回反归一化后的动作所以你正常走流程不会出问题。但如果你自己改代码、自己动手实现推理逻辑这块就特别容易踩坑。在实际项目中我一般会建议不要完全迷信归一化参数是固定的。数据采集的环境光照变化、机械臂安装位置偏移都会影响状态特征的分布。如果换了环境最好重新生成一下归一化参数或者至少做一次快速验证否则模型的表现可能明显下降。3.3 官方配置里diffusion模块的默认参数拆解LeRobot的配置可以通过LeRobotConfig类来加载比如SO-ARM100官方push任务配置文件里diffusion模块的默认参数大致是这样的obs_horizon: 1只观察当前帧pred_horizon: 16一次预测16步action_horizon: 8只执行前8步num_diffusion_steps: 100训练时的扩散步数num_inference_steps: 10推理时的扩散步数n_action_steps: 8实际执行步数一般和action_horizon一致这几个数字不是随便拍的。obs_horizon设成1是因为SO-ARM100这种桌面任务通常不需要历史观测当前状态已经足够决定动作pred_horizon设成16是在“规划长度”和“训练难度”之间的折中action_horizon设成8是为了让模型在执行一段动作后能及时根据最新观察修正轨迹。如果你发现机械臂动作太“急”或者太“拖”优先调这几个值。另外num_inference_steps设成10会让推理快很多但代价是去噪不够充分动作质量可能轻微下降。你可以试着从10改成5、改成20对比一下动作效果和推理耗时找到自己硬件条件下的最优值。4. 实操从配置到训练再到真机部署的完整流程4.1 训练前的数据和环境检查在跑diffusion模块之前有几个检查项能帮你省掉很多无意义的训练时间。第一确认数据集的格式能被LeRobot正确读取。新版LeRobot主推mcap格式旧数据可能是hdf5。建议直接按官方标准重新录制数据或者用官方脚本转换格式不要自己手工改元数据否则特征匹配很容易出问题。第二检查观察特征和动作特征的键名。LeRobot对特征名有约定比如关节状态通常是observation.state动作是action。如果自定义数据集时键名不一致diffusion模块在取数据时会直接报错而且报错信息可能不明显。第三跑一次单batch的正向传播和反向传播确认网络能通。这个步骤很多人会跳过但我强烈建议做一下。你可以写一个简单脚本从数据集里拿一个batch的数据喂给DiffusionPolicy调用compute_loss方法看看loss有没有正常下降。这一步能提前暴露显存不足、tensor维度不匹配等问题比训练到一半再爆显存舒服得多。4.2 启动训练后怎么看loss曲线才是内行LeRobot用Accelerate库来启动训练命令一般是accelerate launch训练脚本会输出loss信息到终端和日志文件。我第一次训练diffusion模块时盯着终端输出看了半天只看到loss在0.5左右波动心里有点没底。后来才发现LeRobot还有一个更直观的评估工具可以实时查看训练曲线和验证效果。diffusion模块的loss曲线跟普通分类任务的loss含义不太一样。它不是越低越好而是要看能不能收敛到一个稳定区间。我刚跑SO-ARM100 push任务时loss从初始的1.2降到0.25左右就基本持平了之后训练到600步也没明显下降。有些人看到loss不降就慌实际上扩散模块的loss本来就不是无限下降的类型关键是降到一个合理的范围并且验证集上的动作质量可接受。有一个比较好的验证习惯训练过程中定期把checkpoint保存下来然后用lerobot.scripts.eval脚本在仿真环境里跑一轮看看动作的成功率。LeRobot支持在仿真环境里直接验证diffusion模块的效果这一步比单纯看loss曲线直观得多。4.3 从训练好的checkpoint到真机执行训练完的产物是一个.pt或.safetensors权重文件加上一个config.json配置。要部署到SO-ARM100真机流程是这样的加载配置、加载权重、实例化DiffusionPolicy、把机械臂的控制循环接上。LeRobot官方提供了一套lerobot命令来简化这个过程但如果你想集成到自己的机器人控制程序里直接调用DiffusionPolicy.select_action这个核心方法是最灵活的。我在自己项目里封装的部署循环大概是这样从机械臂的关节传感器读取当前关节位置组成observation.state。如果策略还依赖图像就同时抓取摄像头画面组成observation.image。把观察数据交给diffusion_policy.select_action(observation)得到动作序列。按action_horizon取出要执行的前若干个动作逐个发送给机械臂控制器执行。循环回到第1步重新读取状态、再次预测。这里有一个特别容易踩的坑select_action的内部会维护一个观察缓冲队列用来拼接多帧历史观察。如果你的控制循环频率不稳定或者中途跳过了某几帧缓冲队列里的数据就乱了动作输出也会跟着乱。LeRobot对这个队列的处理是默认从当前帧开始重新填充不需要手动维护但你得保证每次调用时传入的观察格式一致不能一会传tensor一会传numpy。4.4 实际调参经验我从失败里总结的参数调整思路第一次用diffusion模块跑真机时我遇到的第一个问题是动作迟滞。机械臂执行完一个动作后要停一小会儿才有下一个动作。排查了一圈发现是action_horizon设得太大模型预测了16步但只执行前8步执行完后又重新预测导致每个预测周期之间有一段计算时间。把action_horizon调小到4之后迟滞现象明显好了。第二个问题是动作抖动。关节点在目标位置附近来回震荡像是PID参数没调好但我确认了底层控制没问题。后来发现是scheduler去噪步数太少默认10步去噪不够精细输出动作序列里有微小的高频噪声。把推理时的扩散步数调到15之后抖动明显减轻。第三个问题是运动范围过大机械臂有“甩出去”的感觉。这是因为训练数据和推理时的初始噪声分布不匹配。我的做法是降低归一化输出范围或者在训练数据里删掉那些动作幅度特别大的异常演示数据干净了模型行为自然稳。5. 常见问题与排查技巧我踩过的那些坑希望你不用再踩5.1 数据与环境的坑hdf5文件打不开或者读取缓慢老版本LeRobot用hdf5新版转向mcap。如果你手头只有旧hdf5数据建议用官方脚本转成mcap再做训练否则某些新功能会不兼容。I/O慢的问题可能是连续读文件时没有开缓存批量拷贝到内存里能让训练提速不少。CUDA OOMdiffusion模块的UNet比普通全连接网络大得多默认图像输入加上16步动作预测显存占用轻松超过6GB。如果显存不够优先降低batch size其次降低图像分辨率。我的经验是把图像从256x256降到192x192显存能省下接近一半。训练时loss为NaN大概率是学习率过大或者数据里有异常数值。先检查原始动作数据里有没有NaN再检查归一化时是否出现了除零某个特征的标准差为0。后者有两种情况一种是数据量太少一种是某个关节几乎没有动过。5.2 模型训练与推理的坑模型输出动作总是同一个值这通常说明模型没有学到有效的条件生成观察条件没有有效传入网络。检查一下输入的观察特征是否全部是0或者方差极小常见原因是相机图像没成功加载灰度图被当成全黑图。训练loss正常但真机完全不动这种情况十有八九是状态特征的键名不匹配。模型在训练时用的键是observation.state推理时你传入的dict里键名却写成了obs或者stateLeRobot查不到对应数据就会用默认值填充模型当然不动。解决办法很简单打印一下select_action里实际收到的observation_keys跟训练配置里的feature key对一遍。推理速度太慢机械臂跟不上优先减少num_inference_steps从20降到10速度能快一倍如果还慢就看有没有开启FP16半精度。LeRobot支持在推理时用半精度模型速度提升非常明显但对动作质量有一定影响需要实测权衡。5.3 一套有效的调试优先顺序很多人遇到问题习惯先怀疑模型结构或者学习率但其实排查顺序应该是先确认数据没问题再确认预处理没问题然后确认训练loss合理最后才去动模型参数。我个人调试diffusion模块的优先顺序是这样的用官方demo数据配上官方配置训练一轮确保整个流程能通。换自己的数据不改配置训练一轮对比loss曲线。如果loss明显差很大回到数据本身检查归一化结果、动作分布范围、图像内容。如果loss正常但真机不行进仿真环境反复测试确定是策略问题还是底层控制问题。都排除不了再动pred_horizon、obs_horizon、扩散步数这些超参数。按这个顺序走我几乎没有遇到过解决不了的问题也避免了很多“乱调参导致问题雪上加霜”的情况。6. 经验总结与后续方向最后分享一点个人体会。diffusion模块在LeRobot里确实是效果最稳的默认选择之一但同时也是参数最多、最需要理解的模块。不要一上来就追求改造网络结构先把默认配置跑通跑出一个baseline再一步步去调整解耦。我发现一个提高效率的小技巧把常用的配置文件整理成一个模板每次新任务只改数据路径、特征键、动作维度这几个字段。这样既减少配置错误的概率也让对比实验更有说服力。调参的时候一次只改一个变量不要同时动三四个参数不然根本不知道是哪个改动起了作用。这个模块后续还有很大的扩展空间。比如用更高效的调度器减少推理步数、引入更精细的观察编码器、或者尝试扩散模型与强化学习的结合。我自己接下来的计划是把SO-ARM100默认的push任务换成更复杂的双臂协同任务测试一下diffusion模块在更高维度动作空间里的表现到时候有了结果再写笔记。