
LeRobot 中 Diffusion Policy 的 EMA 权重训练与评估指南完整配置、源码解析与实践【免费下载链接】lerobot LeRobot: Making AI for Robotics more accessible with end-to-end learning项目地址: https://gitcode.com/GitHub_Trending/le/lerobot导读本文围绕 LeRobot 仓库中 Diffusion Policy 官方说明文档 展开系统讲解如何在 LeRobot 中复现 Diffusion PolicyChi et al. 2024参考实现的指数移动平均EMA权重训练与评估流程包括通过lerobot-train开启 EMA、用lerobot-eval评估 EMA 权重、EMA 衰减调度warmup schedule 与常数衰减的配置细节以及 EMA 影子权重在检查点、分布式训练与评估中的底层实现机制。读完本文你将掌握在 LeRobot 上训练 Diffusion Policy 时如何启用并正确评估 EMA 权重并能按需切换为 openpi 风格的常数衰减调度。说明本文以 src/lerobot/policies/diffusion/README.md 为骨架展开文中涉及的源码、配置与测试均可在当前仓库中直接查看验证。Diffusion Policy 在 LeRobot 中的定位从论文到 LeRobot 实现Diffusion Policy 是 Columbia AI 与 HuggingFace 团队联合引入 LeRobot 的代表性策略之一其核心思想是将动作生成建模为条件扩散去噪过程策略不是直接回归动作而是从噪声动作轨迹出发经过多步反向扩散迭代逐步还原出符合观测条件的动作序列。在 LeRobot 中Diffusion Policy 的完整实现位于 src/lerobot/policies/diffusion/由四个文件构成configuration_diffusion.pyDiffusionConfig定义输入输出特征、视觉主干、U-Net、噪声调度器等全部超参数modeling_diffusion.pyDiffusionPolicy、DiffusionModel、DiffusionConditionalUnet1d等模型主体processor_diffusion.py预/后处理管线构造函数init.py模块导出入口。策略类注册名为diffusion因此命令行中通过--policy.typediffusion即可选择该策略见 configuration_diffusion.py 中的PreTrainedConfig.register_subclass(diffusion)与 modeling_diffusion.py 中的name diffusion。原文档的核心诉求复现参考实现的 EMA 行为官方 README 开篇即点明本文主题The reference implementation maintains an exponential moving average (EMA) of the policy weights during training and evaluates the EMA weights.参考实现real-stanford/diffusion_policy在训练期间会对策略权重维护一个指数移动平均副本EMA shadow并在评估阶段使用 EMA 权重而非实时权重。LeRobot 默认关闭 EMA因为需要额外维护一整份参数副本因此若想复现参考实现的行为必须显式开启。训练阶段如何开启并配置 EMA开启 EMA 的最小命令在原文档给出的示例基础上一个最小的 EMA 训练命令为lerobot-train \ --policy.typediffusion \ --ema.enabletrue \ ...其中...表示其余训练参数如--dataset.repo_id、--output_dir等此处省略以便聚焦 EMA 相关配置。--ema.*参数全部由 EMAConfig 数据类定义并通过 train.py 中的 TrainConfig 挂载到训练配置树上。EMA 参数全景表以下是 EMAConfig 定义的完整参数、默认值及含义供配置时对照参考参数默认值含义ema.enablefalse是否在训练时维护 EMA 影子权重。默认关闭因为会额外保留一份完整参数副本ema.decayNone常数衰减系数openpi 风格例如 pi0/pi05 用0.99。设置后绕过下述 warmup 调度每一步都使用该固定衰减不设置则使用 warmup 调度ema.update_after_step0在此步数之前影子权重保持为实时权重的硬拷贝不更新ema.inv_gamma1.0warmup 调度参数见下式ema.power0.75warmup 调度参数见下式ema.min_decay0.0warmup 调度衰减下界clamp 下限ema.max_decay0.9999warmup 调度衰减上界clamp 上限也是最终收敛到的衰减值ema.use_for_evaltrue周期性环境评估时是否使用 EMA 权重离线 eval-loss 评估不受此控制见下文warmup 衰减调度参考实现默认值当不设置ema.decay即使用 warmup 调度时每一步的衰减系数按 diffusersEMAModel的 warmup 公式计算decay_t 1 - (1 t / inv_gamma) ** -power并将结果 clamp 到[min_decay, max_decay]区间内。默认参数inv_gamma1.0、power0.75、min_decay0.0、max_decay0.9999与参考实现保持一致——这正是原文档所述 The EMA decay schedule (--ema.inv_gamma,--ema.power, ...) defaults to the reference implementations values 的含义。该调度在训练早期衰减值很小影子权重几乎实时跟随实时权重随着训练步数增长逐步逼近max_decay0.9999实现先跟随、后平滑的效果。常数衰减对齐 openpi 的 pi0/pi05若希望每一步使用固定衰减而非 warmup 调度例如对齐 openpi 训练 pi0/pi05 时的ema_decay0.99行为只需设置lerobot-train \ --policy.typediffusion \ --ema.enabletrue \ --ema.decay0.99 \ ...从源码看常数衰减的底层实现方式是训练脚本检测到ema.decay非空时将min_decay与max_decay同时固定为该值从而把 warmup 曲线钉在常数上# lerobot_train.py 中的等价逻辑 min_decay cfg.ema.min_decay if cfg.ema.decay is None else cfg.ema.decay max_decay cfg.ema.max_decay if cfg.ema.decay is None else cfg.ema.decay见 lerobot_train.py。正因如此EMAConfig.post_init会强制校验一旦设置了ema.decaymin_decay与max_decay必须保持默认值0.0与0.9999否则抛出 ValueError——两条配置路径互斥避免配置歧义。检查点中的 EMA 权重直接可加载开启 EMA 后每个保存的检查点中会额外携带一份可直接加载的 EMA 权重副本。原文档给出的评估命令即为lerobot-eval --policy.pathoutputs/train/.../checkpoints/last/pretrained_model_ema ...检查点目录结构从 lerobot_train.py 可以看到保存逻辑实时权重保存到checkpoints/step/pretrained_model/目录EMA 权重以_ema后缀的独立目录保存即checkpoints/step/pretrained_model_ema/其中包含完整的模型权重、配置、预处理器与后处理器cfg.save_pretrained、preprocessor.save_pretrained、postprocessor.save_pretrained一并保存同时EMA 影子状态还会以ema_state.pt文件单独存档到training_state/目录下见 lerobot_train.py用于训练中断后的精确恢复断点续训时从该文件恢复影子状态见 lerobot_train.py。因此pretrained_model_ema目录与pretrained_model目录结构完全等价、可直接用--policy.path指向加载——这正是lerobot-eval直接评估 EMA 权重的机制基础。评估 EMA 权重的完整示例结合真实可用参数完整的评估命令形如lerobot-eval \ --policy.pathoutputs/train/act_dt0.1/checkpoints/last/pretrained_model_ema \ --env.type... \ ...--policy.path也可指向 Hub 上的*-ema仓库。若你使用lerobot-push或训练后推送EMA 权重会被推送到独立的repo_id-ema仓库见 lerobot_train.py同样可用于评估。训练内部机制EMA 影子权重如何工作创建与更新流程EMA 影子权重由 diffusers 的EMAModel实例承载在 lerobot_train.py 中完成创建关键逻辑如下前置校验开启 EMA 时若使用 FSDP2/HSDP/CP 等参数分片训练会直接抛出NotImplementedError因为分片场景下参数分散在多个 rank 上无法维护完整影子副本lerobot_train.py同时 EMA 与 PEFT adapter 互斥lerobot_train.py懒加载 diffusers只有开启 EMA 时才导入diffusers.training_utils.EMAModel保证基础训练路径不依赖 diffuserslerobot_train.py仅主进程持有is_main_process()判断后只在主进程创建影子配合 DDP 下各 rank 在梯度同步后权重一致的特性保证安全性见 lerobot_train.py 注释逐优化步更新每完成一个优化器步将实时权重拉入影子见 lerobot_train.py 附近并仅在同步步上执行代码注释注明按同步步而非微批次门控见 lerobot_train.py。评估时如何切换到 EMA 权重lerobot_train.py 中实现了use_ema_for_eval逻辑当 EMA 开启且ema.use_for_evaltrue时周期性环境评估会把 EMA 影子权重临时换入评估用策略模型评估结束后再恢复实时权重。换入换出通过上下文管理器_ema_weights实现lerobot_train.py。值得注意的限制源码注释明确说明见 EMAConfig离线 eval-loss 评估--eval_steps始终使用实时权重。原因是离线评估会在每个 rank 上运行而 EMA 影子只存在于主进程无法在离线评估路径上统一使用 EMA 权重。因此如果你想观察 EMA 权重的真实性能应当通过周期性环境评估或lerobot-eval指向pretrained_model_ema进行。训练日志中的 EMA 信息开启 EMA 后训练启动日志会打印当前使用的衰减配置便于核对是否与预期一致lerobot_train.py常数衰减路径EMA enabled: decay%g (constant), update_after_step%d, use_for_eval%swarmup 调度路径EMA enabled: max_decay%g, inv_gamma%g, power%g, update_after_step%d, use_for_eval%s。端到端实践从训练到评估的完整流程训练配置示例完整命令将上述要点整合一份可复现参考实现行为的完整训练命令如下其余训练参数按需填写lerobot-train \ --policy.typediffusion \ --ema.enabletrue \ --ema.use_for_evaltrue \ --output_diroutputs/train/diffusion_ema \ ...若你的数据集规模较大、希望更平滑的权重平均可调高max_decay若希望影子权重更早开始平滑而非先硬拷贝若干步可保持update_after_step0。评估 EMA 权重训练结束后用指向 EMA 权重目录的方式评估lerobot-eval \ --policy.pathoutputs/train/diffusion_ema/checkpoints/last/pretrained_model_ema \ ...参考实现中的核心默认配置为了让复现参考实现落到实处这里补充 DiffusionConfig 中与 EMA 配合使用的关键默认超参数这些同样来自参考实现可与 EMA 一起原样使用动作分块相关n_obs_steps2观测历史步数、horizon64扩散模型预测的动作序列长度、n_action_steps32每次调用策略实际执行的动作步数并满足n_action_steps horizon - n_obs_steps 1见 modeling_diffusion.py 的说明噪声调度器noise_scheduler_typeDDPM、num_train_timesteps100、beta_schedulesquaredcos_cap_v2、beta_start0.0001、beta_end0.02、prediction_typeepsilon、推理时num_inference_steps默认与训练步数一致见 modeling_diffusion.pyU-Net 结构down_dims(512, 1024, 2048)、kernel_size5、n_groups8、diffusion_step_embed_dim128、use_film_scale_modulationtrue视觉主干vision_backboneresnet18torchvision 预训练权重、spatial_softmax_num_keypoints32、use_separate_rgb_encoder_per_cameratrue优化器与调度器optimizer_lr1e-4、optimizer_betas(0.95, 0.999)、optimizer_weight_decay1e-6、scheduler_namecosine、scheduler_warmup_steps500。注意horizon必须是 U-Net 下采样倍数2 ** len(down_dims)的整数倍DiffusionConfig.__post_init__会在配置校验阶段强制执行configuration_diffusion.py。动手验证用教程脚本跑通 Diffusion Policy仓库提供了两个可直接运行的教程脚本可用于验证上述 EMA 配置的实际效果examples/tutorial/diffusion/diffusion_training_example.py用DiffusionConfigDiffusionPolicy在 LeRobotDataset 上训练完整展示了input_features/output_features推导、delta_timestamps构造依据cfg.observation_delta_indices与cfg.action_delta_indices、make_pre_post_processors构建预/后处理器、policy.forward(batch)计算损失、optimizer.step()更新以及save_pretrained/push_to_hub保存流程。该脚本走的是纯 Python API 训练路径不经过lerobot-train因此不包含 EMA 逻辑——这也反向印证了 EMA 是lerobot-train训练管线的能力examples/tutorial/diffusion/diffusion_using_example.py展示加载训练好的DiffusionPolicy.from_pretrained后在 SO-100 机器人上做闭环推理build_inference_frame构造观测帧、preprocess预处理、model.select_action(obs)选择动作、postprocess反归一化、make_robot_action转换为机器人动作后robot.send_action执行。若你使用lerobot-train训练并开启 EMAcheckpoints/last/pretrained_model_ema目录可直接替代教程中的model_id指向本地路径加载。与其他策略的 EMA 用法对比EMA 机制并非 Diffusion Policy 专属而是 LeRobot 训练管线的通用能力EMAConfig定义于 src/lerobot/configs/default.py。例如 pi0/pi05 的 READMEdocs/source/policy_pi05_README.md同样给出--ema.enabletrue --ema.decay0.99的组合这正是原文档所指的openpi 风格常数衰减。二者的区别仅在于Diffusion Policy 默认走 warmup 调度对齐 Chi et al. 参考实现而 pi0/pi05 官方配置倾向常数衰减对齐 openpi。这一点在 EMAConfig 类注释 中有明确说明。结语与检查清单要在 LeRobot 上完整复现 Diffusion Policy 参考实现的 EMA 训练与评估行为请确认以下几点训练时设置--ema.enabletrue其余调度参数可全部保持默认即参考实现默认值评估时使用--policy.path.../checkpoints/last/pretrained_model_ema指向 EMA 权重副本而非实时权重目录若需对齐 openpi 的 pi0/pi05 训练改用--ema.decay0.99常数衰减注意该选项与min_decay/max_decay互斥离线 eval-loss--eval_steps始终使用实时权重评估 EMA 请依赖周期性环境评估或lerobot-evalEMA 目前不支持 FSDP2/HSDP/CP 分片训练与 PEFT adapter使用前确认训练配置组合断点续训时EMA 影子状态从training_state/ema_state.pt精确恢复无需重新预热。参考与延伸阅读策略说明原文档src/lerobot/policies/diffusion/README.mdEMA 配置定义src/lerobot/configs/default.pyEMA 训练实现src/lerobot/scripts/lerobot_train.py策略配置与模型src/lerobot/policies/diffusion/configuration_diffusion.py、src/lerobot/policies/diffusion/modeling_diffusion.py教程示例examples/tutorial/diffusion/diffusion_training_example.py、examples/tutorial/diffusion/diffusion_using_example.py原始论文引用Diffusion Policy: Visuomotor Policy Learning via Action DiffusionCheng Chi 等The International Journal of Robotics Research, 2024可参见 src/lerobot/policies/diffusion/README.md 中的 BibTeX 条目。提示开启 EMA 会额外保留一份与策略参数等大的权重副本显存占用会相应增加源码注释明确指出这是默认关闭的原因见 EMAConfig 类注释。在显存受限的环境下请权衡 EMA 带来的评估稳定性提升与额外的显存开销。【免费下载链接】lerobot LeRobot: Making AI for Robotics more accessible with end-to-end learning项目地址: https://gitcode.com/GitHub_Trending/le/lerobot创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考