我做机器人动作模仿训练有一段时间了lerobot_train 这个训练脚本我几乎每天都要打交道。说实话大部分人在复现动作模仿项目时注意力都放在 policy 结构、数据增强、视觉编码器上scheduler 常常被认为是“配置里一个可有可无的字段”。可如果你真的把训练日志里的 learning rate 曲线画出来就会发现收敛速度、最终精度、甚至偶发的 loss 变成 nan都和 scheduler 的具体实现有直接关系。这篇文章想把 lerobot_train 里 scheduler 从配置解析、类定义、步进时机到恢复训练这整条链路讲清楚。我尽量按实际工程里会遇到的方式去讲而不是只给一个“这是学习率调度器”的概念。适合刚开始接触该类框架训练脚本的人也适合已经能跑通训练、但想弄明白 scheduler 到底在哪个环节生效、为什么改动一个参数会导致完全不同的训练曲线的同学。1. 先搞清楚 scheduler 在 lerobot_train 里管的是什么1.1 学习率为什么不能一直固定在一个值很多刚接触训练的人会有一个直觉学习率越大收敛越快学习率越小越稳定。于是干脆把 lr 设成一个固定值比如 1e-4一直训练到底。这种做法不是完全不能训练但大多数时候你会发现两个问题训练刚开始时 loss 很高固定 1e-4 偏大模型参数在一个不稳定的初始位置来回震荡训练到后期 loss 已经很低了固定 1e-4 又偏大导致 loss 始终在一个小范围内下不去或者甚至出现 loss 突然冲高又回落的“抽风”现象。scheduler 存在的意义就是把学习率这个单一超参变成一条随时间变化的曲线。前期用一个较小或逐渐上升的学习率让模型平稳起步中期用较大学习率快速优化后期再逐步衰减让参数落进一个比较平滑的局部最优区域。这个思路在大多数优化器上都适用不管你是 AdamW、SGD 还是 AdaFactor。我用一个类比学习率很像你开车时的油门。刚上赛道时你不可能直接踩死油门因为轮胎还没热、路面情况还没摸清所以要循序渐进地加油等到路况清楚了可以大油门冲一段到最后一个弯道又必须减速否则一定会冲出赛道。scheduler 就是那个帮你规划“什么时候加油、什么时候减速”的脚本。1.2 lerobot_train 整体链路里 scheduler 处于哪个位置在 lerobot_train 这种训练脚本里一个典型训练循环大概是这样组织的加载配置文件和数据集构建 policy 模型构建 optimizer根据配置构建 scheduler进入 epoch 和 batch 循环计算 loss、反向传播、更新参数在合适的时机调用 scheduler.step()周期性保存 checkpointcheckpoint 里同时包含模型、optimizer、scheduler 的状态。scheduler 本身不参与梯度计算它只做一件事告诉 optimizer 当前应该用多大的学习率。PyTorch 里的实现方式是通过修改 optimizer.param_groups 中每一项的 lr 字段。所以你可以直接把 scheduler 理解为 optimizer 的学习率控制器没有 optimizer 就没有 scheduler。在 lerobot_train 中配置文件里通常会看到类似这样的字段training: optimizer: name: adamw lr: 1.0e-4 weight_decay: 1.0e-4 scheduler: name: cosine_annealing_warm_restarts num_cycles: 4 warmup_steps: 500 max_steps: 100000这段配置的意思很直白用 AdamW 作为优化器基础学习率 1e-4再用一个带 warmup 的余弦退火重启 scheduler总共训练 10 万步其中前 500 步做线性热身之后按 4 个周期做余弦衰减。n 这个warmup_steps在很多人看来只是热身多少步但这个参数直接决定了训练初期模型是否稳定其实际影响比想象中要大。2. scheduler 的初始化与配置加载拆解2.1 配置文件如何变成内存里的 scheduler 对象我在自己写的训练脚本里通常会用一个统一的配置类来管理这些字段。lerobot_train 这一类脚本也是类似思路配置字段先被读取成一个字典或 dataclass然后根据 scheduler.name 去选择对应的 scheduler 类再传入 optimizer 和若干调度参数。简化后的创建逻辑大概是这样from dataclasses import dataclass dataclass class TrainingConfig: max_steps: int 100000 warmup_steps: int 500 lr: float 1.0e-4 def build_optimizer_and_scheduler(model, training_cfg): optimizer torch.optim.AdamW( model.parameters(), lrtraining_cfg.lr, weight_decay1.0e-4, ) scheduler_cfg { name: cosine_annealing_warm_restarts, num_cycles: 4, warmup_steps: training_cfg.warmup_steps, max_steps: training_cfg.max_steps, } scheduler build_scheduler(optimizer, scheduler_cfg) return optimizer, scheduler def build_scheduler(optimizer, scheduler_cfg): name scheduler_cfg[name] if name cosine_annealing_warm_restarts: return CosineAnnealingWarmRestartsWithWarmup( optimizer, max_stepsscheduler_cfg[max_steps], warmup_stepsscheduler_cfg[warmup_steps], num_cyclesscheduler_cfg[num_cycles], ) elif name constant: return ConstantScheduler(optimizer) else: raise ValueError(fUnsupported scheduler: {name})这里比较关键的一点是scheduler 的创建必须拿到 optimizer 对象因为 scheduler 需要读取 optimizer.param_groups 里的初始学习率作为自己的基准。也就是说如果 optimizer 的 lr 配置写错了scheduler 再正确也没用它只会基于那个错误的初始 lr 去生成后续曲线。我在实际工程中发现很多人把 optimizer 的 lr 和 scheduler 的最大 lr 混淆。scheduler 内部通常会保存一份 base_lr这个 base_lr 就是创建 scheduler 时 optimizer 的初始 lr。如果你的自定义 scheduler 里硬编码了一个 lr那配置文件的 lr 字段就失效了这个问题排查起来非常隐蔽。2.2 scheduler 类型怎么选lerobot_train 相关的项目里我见过比较多的 scheduler 有下面几种scheduler 名称学习率趋势适合场景constant全程不变简单任务、短训练、需要精确复现实验cosine_annealing从峰值余弦衰减到 0较长训练任务收敛稳定性好cosine_annealing_warm_restarts余弦衰减后周期性重启动作模仿、扩散策略这类容易陷入局部最优的训练onecycle先升后降并伴随动量调整快速训练、需要在有限步数内达到较好效果reduce_on_plateau指标不降时自动降低 lr需要根据验证集指标动态调整的任务在动作模仿类任务里我比较推荐带 warmup 的余弦退火或余弦退火重启。原因很简单动作模仿任务往往需要在一个固定数据集上反复学模型很容易快速记住训练集里的样本如果不把学习率降下来后期就会出现 loss 在低位震荡、验证效果却变差的情况。加周期重启还能让模型在退火之后跳过一些比较差的局部区域实测下来效果会更稳。这里也顺带提一个思路如果你在跑一个新任务完全不确定该选什么 scheduler最简单的起点是 constant warmup。先跑通流程画出 loss 曲线再根据曲线形状决定是否换余弦退火。不要一上来就在 scheduler 上堆一大堆参数否则出问题时很难判断是 scheduler 的问题还是模型结构的问题。3. scheduler 核心实现warmup 和余弦退火到底怎么算3.1 warmup 阶段的线性升温原理带 warmup 的 scheduler第一步要处理的是从第 0 步到 warmup_steps 结束这段区间。绝大多数实现会采用线性升温也就是当前学习率等于基础学习率乘以当前进度lr base_lr * (current_step / warmup_steps)当 current_step 等于 1 时学习率大概是 base_lr 的千分之一当 current_step 接近 warmup_steps 时学习率逼近 base_lr。这样做的理由是训练最开始模型参数是随机初始化状态梯度方向往往包含大量噪声如果一上来就直接用大学习率更新参数模型很可能会被推到一片非常糟糕的区域再花很多步都回不来。我见过有的实现会对 warmup 阶段的进度做开方或者余弦插值但线性升温已经足够好用。还有一点要注意warmup_steps 不应该占总训练步数的比例太大。如果你训练 10 万步warmup 设 2 万步那等于前 20% 的训练时间都在低学习率下慢慢爬总步数可能不够模型收敛。我常用的经验是 warmup 占 max_steps 的 0.5% 到 5%具体还要看数据集大小和 batch size。3.2 余弦退火阶段的核心公式与周期控制warmup 结束之后真正的主角是余弦退火。基础版本的计算方式是将当前进度映射到 0 到 1 之间然后计算余弦值学习率从峰值逐渐衰减到接近 0。标准公式大致是lr base_lr / 2 * (1 cos(pi * progress))progress 从 0 到 1 时cos 从 1 变到 -1所以 lr 从 base_lr 变到 0。这个曲线的好处是前期衰减比较慢模型有足够时间在大学习率下探索后期衰减越来越快能帮助模型进入更精细的参数区域。带 warm restart 的版本会在此基础上把整个训练区间切成多个周期每个周期内都做一次完整的余弦退火。比如 num_cycles 设成 4那么 10 万步里大约每 2.5 万步经历一次从峰值到 0 的退火然后学习率重新弹回峰值开始下一个周期。这种“重启”的意义是当一个周期结束时模型被推到一个局部低点立刻增大学习率模型就会跳出这个低点去寻找另一个可能更好的区域。我这里写一个简化的实现逻辑和我在 lerobot_train 里实际用到的版本一致import math class CosineAnnealingWarmRestartsWithWarmup: def __init__(self, optimizer, max_steps, warmup_steps, num_cycles1): self.optimizer optimizer self.base_lr optimizer.param_groups[0][lr] self.max_steps max_steps self.warmup_steps warmup_steps self.num_cycles num_cycles self.current_step 0 def get_lr(self): step self.current_step if step self.warmup_steps: return self.base_lr * step / max(1, self.warmup_steps) after_warmup step - self.warmup_steps total_anneal_steps self.max_steps - self.warmup_steps cycle_len total_anneal_steps / self.num_cycles pos_in_cycle after_warmup % cycle_len progress pos_in_cycle / cycle_len return self.base_lr / 2 * (1 math.cos(math.pi * progress)) def step(self): self.current_step 1 lr self.get_lr() for param_group in self.optimizer.param_groups: param_group[lr] lr这段代码里有个容易被忽视的细节after_warmup % cycle_len利用取模来模拟“多个周期”。比如 num_cycles4每个周期长度是 2.5 万步那么第 3 万步时 after_warmup 为 2.95 万步取模之后约等于 4500 步对应的 progress 是 0.18所以学习率还在一个比较高的位置。这比手动维护“当前周期下标”要简单得多也能避免重启时出现 off-by-one 错误。我在调试时发现周期长度如果不是整数取模逻辑可能会让最后一段曲线少几步导致最后一个周期无法衰减到一个很小的值。这时可以在配置里把 max_steps 和 num_cycles 的除法结果向上取整或者在构造周期长度时用math.ceil。这个细节对最终精度有一点影响但不致命。3.3 与 optimizer.param_groups 的交互机制上面代码里每调用一次step()就会把 optimizer.param_groups 中所有 group 的 lr 改成新值。这里有一个容易踩的坑如果模型里不同参数组设置了不同学习率比如 encoder 用 1e-4policy head 用 1e-3而 scheduler 的 base_lr 只取了第一个 param_group 的 lr那么所有组都会被同步成同一个值导致你的分组学习率设置失效。解决方式有两种。第一种是 scheduler 内部保存每个 group 的初始 lr 数组然后在计算时按 group 分别处理。第二种更常见在训练配置里不用多层 param_group所有参数都用同一个学习率scheduler 的代码就简单很多。如果你确实需要分组学习率建议在 scheduler 的__init__里记录所有初始 lrself.base_lrs [group[lr] for group in optimizer.param_groups]计算时和 base_lr 相乘再写回对应的 param_group。这样才能保证分组设置不会被覆盖。4. 训练循环里 scheduler 的调用位置与步进节奏4.1 到底是在每个 batch 步进还是每个 epoch 步进很多初学者会在配置里看到 max_steps这个字段指的通常是优化器更新次数而不是 epoch 数。所以 scheduler.step() 最自然的调用位置是在 optimizer.step() 之后、下一个 batch 开始之前。这样一来学习率的变化是跟着参数更新走的而不是跟着数据遍历轮数走。我在 lerobot_train 里常用的循环方式如下for step, batch in enumerate(train_dataloader): optimizer.zero_grad() loss policy.compute_loss(batch) loss.backward() optimizer.step() if scheduler is not None: scheduler.step() if step % log_interval 0: current_lr optimizer.param_groups[0][lr] print(fstep{step}, lr{current_lr:.6f}, loss{loss.item():.4f})把 scheduler.step() 放在 optimizer.step() 之后是为了保证当前步的训练用得还是上一步结束时的学习率而下一步训练时学习的率才更新。如果你把顺序反过来先 scheduler.step() 再 optimizer.step()相当于第一次更新参数时直接跳过了 warmup 前几步学习率曲线会比预期提前一个 step虽然只差一步但在严格复现实验时会造成轻微误差。至于“每个 epoch 步进”的做法在图像分类任务里比较常见但在动作模仿和扩散策略这类任务里我建议使用按 step 的方式。因为一个 epoch 内样本数量和 batch size 可能经常变动按 epoch 步进会让学习率曲线受到数据加载顺序影响不够稳定。按 step 步进则完全可控。4.2 checkpoint 保存和恢复时必须带上 scheduler 状态我在实际工程中踩过一个大坑训练到第 8 万步时因为机器需要重启我恢复了 checkpoint 继续训练结果 loss 不但没有继续下降反而回到了训练初期的水平。排查半天发现checkpoint 里只保存了 model 和 optimizer 的 state_dictscheduler 的 current_step 被重置为 0于是 warmup 阶段又重新开始了整个学习率曲线被拉回起点。正确的做法是在保存时同时保存 scheduler 状态并在恢复时加载checkpoint { model: model.state_dict(), optimizer: optimizer.state_dict(), scheduler: scheduler.state_dict(), step: step, } torch.save(checkpoint, checkpoint.pt) # 恢复时 checkpoint torch.load(checkpoint.pt) model.load_state_dict(checkpoint[model]) optimizer.load_state_dict(checkpoint[optimizer]) scheduler.load_state_dict(checkpoint[scheduler]) step checkpoint[step]如果你的 scheduler 是自定义类state_dict 里至少应该包含 current_step、max_steps、warmup_steps、num_cycles、base_lr 这些关键字段。恢复后如果发现 lr 曲线不对优先去检查 current_step 是否被正确恢复了这是最简单的排查方向。4.3 多卡并行训练时 scheduler 的同步问题在单卡训练时 scheduler 的一切都很简单。但一旦切到 DistributedDataParallel不同 rank 上都会有独立的 optimizer 和 scheduler。如果每个 rank 都从 dataloader 里自己取数据理论上每个 rank 的 step 数量相同scheduler 的 current_step 也会同步。问题往往出现在“每个 rank 的 batch 数量不完全一致”或“有 rank 提前退出”的情况。我在多卡训练时一般约定scheduler.step() 在每张卡上都正常调用不能只在 rank 0 上调用否则其他 rank 的学习率和 rank 0 不一致模型参数更新就会出现差异。同时保存 checkpoint 时只让 rank 0 执行保存之前把 optimizer 和 scheduler 的 state_dict 都收集到 rank 0 上。由于每个 rank 训练步数基本一致state_dict 内容也一致所以不会有什么问题。还需要注意一点如果把 batch size 调大等价于每个 step 看到更多样本训练总步数通常应该减少。这时如果你只调了 batch size 而没调 max_stepsscheduler 在同样的总步数里走完退火而模型实际看到的样本量变多了学习率衰减节奏会变快。这不算 bug但是一个经常被忽略的联动因素。5. 常见问题与排查技巧实录5.1 loss 不下降或者直接变成 nan遇到这种情况第一反应不要去看模型结构先确认当前真实学习率是多少。我在训练脚本里有时会漏掉一行current_lr optimizer.param_groups[0][lr]结果日志里根本没记录学习率出了问题只能靠猜。加上学习率日志之后先看前三五千步学习率是否从很小逐渐涨到设定峰值。如果没有说明 scheduler.step() 可能没被调用或者 warmup 逻辑写错了、current_step 一直停在 0。学习率太大导致 nan 的表现是loss 在前几百步内忽大忽小然后突然出现 inf 或 nan。这时候可以看看峰值学习率是多少如果是 1e-3 甚至更高试着把 lr 降到 1e-4 或者 3e-5很多情况下问题就消失了。需要明白 scheduler 只是把学习率最大值控制在 base_lr 上如果你的 base_lr 本身设置过高warmup 再长也救不回来。5.2 学习率曲线后期突然“跳回去”这种问题几乎都出在带 warm restart 的 scheduler 上。如果 num_cycles 大于 1每个周期结束后学习率都会弹回到峰值。有些任务确实适合这种重启但也有任务会因为重启后学习率瞬间变大导致训练不稳定。我的经验是如果重启后 loss 剧烈抖动可以把重启次数从 4 降到 1也就是退化成普通余弦退火先观察效果。还有一个小问题如果你在训练中手动调整了 optimizer 的 lr比如写了optimizer.param_groups[0][lr] 1e-5scheduler 下一次 step 时会直接覆盖这个手动调整值。所以不要试图在训练循环里手动修改 lr要改就改 scheduler 内部的参数否则会出现“你改了但没生效”的错觉。5.3 恢复训练后学习率不对恢复 checkpoint 后学习率不对通常有三种可能。第一种是 checkpoint 里根本没有 scheduler state_dict第二种是 current_step 从 0 重新开始第三种是 scheduler 类内部的__init__参数和保存时不一致导致 state_dict 加载时报错或静默失败。我自己在写自定义 scheduler 时会在 state_dict 里保存一个version字段。以后如果改了调度逻辑加载时看到 version 不一致就主动报警避免静默加载到错误状态。这是一个低成本但很有效的自我保护措施。5.4 手写公式跟 PyTorch 内置调度器对不上如果你使用 PyTorch 内置的 CosineAnnealingWarmRestarts 或 OneCycleLR建议先手动打印前 50 步的学习率跟自己的计算公式对比一下。内置调度器的 T_0、T_mul、epoch 参数含义有点绕特别是 T_mul 默认是 2表示每重启一次周期长度翻倍。如果你配置里写了 num_cycles4但底层用的是 T_mul2那么实际周期长度不是均匀的学习率曲线会是一个逐渐拉长的过程。如果项目对可复现性要求很高我倾向于直接用自定义 scheduler因为它的逻辑完全透明日志里每一步对应什么公式都清清楚楚。内置调度器虽然方便但在配置含义上需要读更多文档。6. 一些个人经验和调试小技巧最后分享一个我调 scheduler 时常用的方法把学习率变化直接画成图片和 loss 曲线放在同一个时间轴上。很多人只习惯于记住第一个和最后一个学习率但中间曲线的形状对训练结果影响很大。比如同样是 10 万步训练一个用 constant一个用 cosine最终验证精度可能差 3 到 5 个点。如果不画图很难发现“原来学习率在中期就已经衰减到几乎为 0”这种问题。另外一个技巧是在 warmup 结束的那个 step 前后手动打印一次 lr 和 loss。如果 loss 在 warmup 结束时出现尖峰说明峰值学习率还是太大可以适当降 lr 或者拉长 warmup_steps如果 loss 没有任何波动说明模型容量或数据质量才是主要瓶颈不要继续乱调 scheduler。我建议每个训练项目都在配置里保留一组“scheduler baseline”参数比如 constant warmup_steps0作为对照组。当你改了模型结构或数据后不要只盯训练结果先用这组 baseline 跑一遍看看到底是模型本身变强了还是 scheduler 的功劳。这种做法能帮你更客观地判断每次改动是否真的有效。如果你准备在现有脚本里加一个新的 scheduler从最简单的类开始写确保 state_dict 可保存、可恢复、可打印。把那几个 step 细节处理好比引入一堆复杂配置要有用得多。
