做机器人模仿学习训练的朋友应该都有类似体验:模型结构、数据采集、回放buffer都搞得明明白白,最后训练一跑起来,loss曲线和预期差得十万八千里,然后开始怀疑是网络设计问题、数据问题,甚至怀疑人生。我在复现某套开源机器人操作训练流程时也踩过这个坑,后来才意识到,问题很可能出在一个平时不太起眼的组件——scheduler。
LeRobot这套训练框架里,lerobot_train是入口,负责拉起模型、数据、优化器、scheduler整套流程。scheduler虽然只在每个step结束做一件事——调整学习率,但它直接影响模型能不能收敛、收敛得快不快、最终精度高不高。很多介绍资料都是一句话带过,可实际去看代码的时候,里面的细节多得惊人。这篇文章我把lerobot_train里scheduler的具体实现从头到尾拆一遍,包含配置声明在哪、怎么被加载、训练循环里如何生效,以及我在实际调参时踩过的几个坑。
1. scheduler在整个LeRobot训练流程中的位置与配置入口
1.1 从启动脚本追踪scheduler的加载路径
先说结论:lerobot_train里的scheduler不是训练脚本临时造的,而是通过配置系统定义好的。整套框架用的是一套基于配置文件的训练启动方式,启动时会加载一个.yaml配置,里面把所有训练相关参数全部串起来,包括模型名、数据集路径、batch size、学习率、以及我们关心的scheduler参数。
启动流程大致是这样的:先解析命令行参数并识别对应任务的配置文件路径,然后加载配置和预训练权重,再构建数据集和DataLoader,接着构建模型、优化器,最后才轮到scheduler。这一步顺序很关键,因为scheduler的构建依赖优化器和训练总步数,而训练总步数又依赖数据集大小和批大小。
我去翻源码时发现,scheduler相关配置通常在配置文件的optimizer或者training字段下面,不同版本Key位置略有调整,但核心没有变。以某个常见配置为例:
optimizer: lr: 1e-4 betas: [0.95, 0.999] eps: 1e-8 weight_decay: 1e-6 scheduler: name: cosine_with_warmup num_warmup_steps: 500 num_training_steps: -1 num_cycles: 0.5这里name指定了调度器类型,num_warmup_steps是预热步数,num_training_steps设置为-1表示训练总步数到后面根据数据集和epoch数推算,num_cycles是余弦退火的周期比例。
1.2 为什么模仿学习训练偏爱cosine_with_warmup
我第一次看到这个配置时有个疑问:为什么不用最经典的StepLR或ExponentialLR,偏偏要用cosine_with_warmup?后来才发现,机器人操作领域的模仿学习任务,普遍训练步数不多,而且数据集规模相对小,学习率稍微调大一点就容易震荡,调小一点又收敛太慢。余弦退火的特点是前期保持较高学习率快速逼近较好区域,后期通过余弦曲线逐步下降到接近0,相当于让模型在后期做精细打磨。
加warmup的原因很朴素:训练刚开始时优化器里的动量估计还是乱的,如果一上来就按初始学习率大步更新,可能直接把模型参数推到不理想的位置。尤其是LeRobot里常用预训练的视觉骨干网络,刚加载时骨干层的梯度分布和预测头差异很大,一小段warmup让学习率从0线性爬到目标值,能有效避免初期loss爆发。
1.3 scheduler和优化器之间的顺序约束
scheduler在构建时必须拿到优化器实例,因为它的core逻辑就是修改优化器内部的学习率参数。LeRobot里构建顺序是固定写死的:先构建模型,然后构建优化器,接着把优化器传给scheduler工厂函数,最后再交给训练循环使用。
这个顺序也有一个隐蔽的坑:如果提前用torch.optim.lr_scheduler改了优化器的param_groups,后续再使用框架的scheduler工厂函数时,可能会因为初始化方式不同而丢掉之前的学习率设置。所以千万不要自己额外套一层torch.optim.lr_scheduler.CosineAnnealingLR之类的调度器,要么用框架内置的,要么就自己完全接管欧拉器,二者只能选一个。
2. get_scheduler构造链条:从配置字段到学习率曲线的关键代码拆解
2.1 训练脚本里实际调用的是哪个工厂函数
LeRobot代码里,scheduler的构建通常封装成一个独立的工厂函数,有些版本直接调用transformers库的get_scheduler。这个函数可以理解为一个统一的调度器开关,输入调度器名字、优化器、预热步数、总训练步数,返回一个torch的调度器对象。
代码层看起来是这样:
from transformers import get_scheduler # cfg.optimizer.scheduler 对应上面配置文件的 scheduler 子字段 scheduler = get_scheduler( name=cfg.optimizer.scheduler.name, optimizer=optimizer, num_warmup_steps=cfg.optimizer.scheduler.num_warmup_steps, num_training_steps=training_steps, )注意num_training_steps这里传进去的是前面计算出的真实训练步数。源码里当配置中设置成-1时,框架会自己根据总数据量、batch size、gradient accumulation步数和epoch数推算。推算逻辑大致是:
steps_per_epoch = len(dataloader) // gradient_accumulation_steps training_steps = steps_per_epoch * num_epochs为什么除以gradient_accumulation_steps?因为有效更新次数才等于真实学习率调整次数,我后面会展开说。
2.2 name字符串映射到具体调度策略
get_scheduler内部维护了一张调度器名称和实现类的映射表。cosine_with_warmup对应的就是带预热的余弦退火;linear对应线性衰减;constant_with_warmup是预热之后保持恒定;还有polynomial、inverse_sqrt、cosine_with_min_lr等多种选择。
这张表格我经常参考:
| name | 学习率走势 | 典型使用场景 |
|---|---|---|
| linear | 从峰值线性降到0 | 通用迁移学习 |
| cosine_with_warmup | 预热后余弦曲线降到0 | 模仿学习、微调大模型 |
| cosine_with_min_lr | 降到min_lr但不归零 | 训练后期仍想保留一定探索 |
| constant_with_warmup | 预热后恒定不变 | 任务简单、步数少 |
| polynomial | 幂次衰减 | 图像生成类任务 |
我实测下来LeRobot的默认配置使用cosine_with_warmup在ALOHA、单臂推杆这类任务上表现最稳。余弦衰减的好处是收敛末期控制得好,模型不会因为学习率太大在最优附近来回抖动。
2.3 num_cycles到底是怎么影响学习率曲线的
num_cycles这个参数容易被忽略,但它决定了余弦曲线的完整周期数。默认0.5时,学习率从峰值对应cos(0),到训练结束时对应cos(pi),正好走半个余弦波。也就是说学习率单调地从峰值下降到接近0,没有回升阶段。
如果你把num_cycles改成1,训练过程就会经历"降-升-降"好几轮,模拟循环学习率的效果。很多人在模仿学习任务上试过增大num_cycles,结果不理想,原因是模仿学习的任务复杂度远低于自监督预训练,不需要反复逃离局部极值,学习率中途回升反而容易破坏后期已经收敛好的动作分布。
2.4 schedule_free优化器对scheduler的影响
LeRobot较新的版本支持schedule_free优化器,这类优化器内置了类似学习率退火的功能,和传统scheduler叠加使用会有冲突。如果选了schedule_free,配置里的scheduler最好设置为空或者constant,否则会出现学习率被调度器强制压低,但优化器的内置退火逻辑又在叠加一套,训练动态变得很难解释。
我记得在某一次对比实验中,同一条数据,把scheduler从cosine_with_warmup换成constant,配合schedule_free优化器,收敛速度反而提升了。但是默认的AdamW加cosine_with_warmup组合仍然是最稳妥的,不建议新手一上来就折腾这个。
3. 训练循环里scheduler的更新时机与分布式环境下的差异
3.1 每个step结束后到底该谁先走
scheduler的更新时机看起来简单——每个step结束之后调用scheduler.step()。但LeRobot里有一个细节容易被忽略:优化器的optimizer.step()和scheduler的scheduler.step()不是紧紧挨着的,中间可能隔着梯度清零、梯度裁剪、EMA更新等操作。不同版本训练脚本的顺序有差异,但大体是:算梯度、反向传播、梯度裁剪、优化器更新、scheduler更新、EMA更新。
顺序的核心原则是:scheduler应该紧跟在optimizer.step之后,确保本轮使用的是当前步的学习率。如果反了,相当于当前步用了下一步的学习率,行为会偏移一点,虽然不炸,但复现时会发现曲线和别人对不上。
3.2 accelerate与scheduler底层封装
LeRobot依赖Hugging Face的accelerate库做分布式训练和混合精度。accelerate提供了一套Accelerator对象,其中有个scheduler输入参数。LeRobot会把上面构建好的scheduler传给Accelerator,之后训练循环里的scheduler.step()实际上是被accelerate接管并分发的。
如果我在自己写训练脚本时用accelerate.prepare(),优化器和scheduler会被包装成分布式版本。此时如果再用scheduler.step(),它会自动处理主进程和其他进程的同步问题。这也是一个容易踩坑的点:自己单独创建scheduler但不通过accelerate.prepare传给Accelerator,分布式训练时每个进程的学习率调整步数会不同步,训练曲线就会开始漂。
3.3 梯度累积时scheduler的更新次数必须除以累积步数
梯度累积逻辑让问题更隐蔽。假设batch size是64,机器显存只能跑batch size 8,那么可以累积8次梯度再更新一次参数。训练循环里通常这样写:
for batch in dataloader: loss = model(batch) loss.backward() if (step + 1) % gradient_accumulation_steps == 0: optimizer.step() optimizer.zero_grad() scheduler.step()这里scheduler.step()只在参数更新时调用一次,而不是每个micro-batch都调。如果忘了做这个除法,总训练步数会变成micro-batch数 × epoch数,结果warmup还没做完,学习率已经开始下降了,整个训练过程的学习率曲线完全变样。
LeRobot里处理这个的方式就是前面提到的,计算训练总步数时除以gradient_accumulation_steps,保证scheduler的更新节奏和优化器保持一致。
3.4 checkpoint恢复训练时scheduler状态要一起恢复
模拟学习训练周期长,跑了一半断电或者OOM是常事。LeRobot的checkpoint保存策略通常除了模型权重、优化器状态,还会保存当前epoch和scheduler状态。但有些版本只保存到优化器,不保存scheduler的last_epoch,恢复训练时会从零开始重新warmup,结果就是学习率突然从0重新爬,前面的训练进度直接作废了一截。
我的习惯是恢复训练后立刻打印一下当前学习率,确认scheduler的步数和日志里的step对得上。如果对不上,就到checkpoint文件里找scheduler_state_dict,手动加载回去。这个小动作能避免掉一半的"恢复训练后loss异常"问题。
4. 调参实战:warmup、num_cycles、步数边界那些容易翻车的细节
4.1 warmup步数怎么选才靠谱
不小比例的人直接用默认的500步,结果小数据集上整个训练才两三千步,warmup占了四分之一,训练前期模型基本在原地踏步。我用的经验规则是:warmup步数占总训练步数的5%-10%比较合理。如果总步数只有2000,那100到200步的warmup就够用了。
注意,如果用预训练骨干网络,warmup可以略微取高一点,比如15%,给骨干网络一个热身时间;如果从头训练模型且数据集小,warmup建议更短,不然模型迟迟进入不了正常学习状态。
4.2 学习率降到底之后模型还能动吗
cosine_with_warmup训练到尾声时学习率趋近于0,模型参数几乎不再更新。有些任务训练末期的动作精度反而变差,不是scheduler的问题,而是早期模型已经收敛到了某个位置,后期学习率太小无法跳出,相当于被锁死了。
解决办法两个思路:一是把name换成cosine_with_min_lr,设置一个极小的下限比如1e-6,让模型保留微调能力;二是训练结束后用最后一个ckpt继续用小学习率微调几个epoch。我在实际操作中偏好前者,因为改动小,训练流程不用断。
4.3 多任务训练和单任务训练要不要不同scheduler
LeRobot有时做多任务策略,一个模型同时学多个动作。多任务数据集大,训练总步数多,warmup固定500步影响很小。但多任务场景下模型更容易出现学习率后期过大导致的灾难性遗忘,所以可以把余弦改成cosine_with_min_lr,并且把num_cycles调小到0.3,让学习率更早降到较低区间,减少对之前学过的动作分布的冲击。
单一任务场景反而可以保持默认的0.5。因为单任务目标一致,学习率快速下降有利于收敛到稳定的动作流形。
4.4 我认为比较稳的一套初始配置
给一套我验证过多次的初始配置:
optimizer: lr: 1e-4 weight_decay: 1e-6 scheduler: name: cosine_with_warmup num_warmup_steps: 200 num_training_steps: -1 num_cycles: 0.5训练总步数2000以内,把预热调成100。总步数过万时,预热可以放到500到800。learning rate这块,1e-4是模仿学习任务的比较稳妥起点,如果模型收敛很慢且loss平稳,可以往上提到2e-4,但要配合同样的warmup比例,否则刚开始几步loss可能出现尖峰。
4.5 如何快速判断scheduler配置是否生效
我在调试时最常用的诊断手段就是打印学习率曲线。训练每跑完一个epoch,把optimizer.param_groups[0]['lr']存到日志里,然后画出来。如果曲线长这样:从0先线性爬到峰值,然后走半个余弦波降到接近0,说明scheduler完全生效。如果曲线一开始就直接是峰值,说明warmup没生效;如果训练还不到一半学习率就到0了,说明总训练步数计算出了问题。
还有一种情况是学习率根本不变,这就要检查是不是选了constant类型,或者代码里没调用scheduler.step()。
5. 当默认scheduler不够用时:如何实现自定义调度策略
5.1 基于torch的原生自定义scheduler写法
有时候余弦退火确实不满足需求。比如动作策略训练我们希望前期学习率大跨步朝数据分布靠拢,中期保持一段恒定,后期再慢慢降,这种三段式策略就需要自己写一个scheduler。
torch的自定义scheduler实现方式很简单,继承torch.optim.lr_scheduler.LRScheduler,重写get_lr方法:
from torch.optim.lr_scheduler import LRScheduler class ThreePhaseScheduler(LRScheduler): def __init__(self, optimizer, warmup_steps, hold_steps, decay_steps, min_lr_ratio=0.05): self.warmup_steps = warmup_steps self.hold_steps = hold_steps self.decay_steps = decay_steps self.min_lr_ratio = min_lr_ratio self.base_lr = optimizer.defaults["lr"] super().__init__(optimizer) def get_lr(self): step = self.last_epoch if step < self.warmup_steps: factor = step / max(1, self.warmup_steps) elif step < self.warmup_steps + self.hold_steps: factor = 1.0 else: progress = min(1.0, (step - self.warmup_steps - self.hold_steps) / max(1, self.decay_steps)) factor = 1.0 - (1.0 - self.min_lr_ratio) * progress return [self.base_lr * factor for _ in self.optimizer.param_groups]然后把它接到训练循环里,替换默认scheduler,逻辑上不需要动其他代码。
5.2 在LeRobot里替换自定义scheduler的接入方式
LeRobot里想要用自定义scheduler,最简单的思路是绕过框架默认的scheduler构建逻辑,在训练循环里自己持有scheduler实例。训练循环里会预留一个scheduler.step()钩子,只要在optimizer构造完之后,自己创建scheduler对象,并把原来的scheduler构建部分置空或跳过,就能无缝替换。
需要注意的部分:一定要把自定义scheduler也能随checkpoint保存和恢复。LeRobot的保存逻辑看的是传入的scheduler对象是否有state_dict方法,torch自带的LRScheduler都有,所以继承LRScheduler写没问题。恢复时直接scheduler.load_state_dict(saved['scheduler'])。
5.3 关于scheduler设计的一个反直觉结论
我以前总觉得学习率越大前期训练越快,直到有一次对比实验让我悔不当初。同样的任务,把峰值学习率从1e-4改成2e-4,前期loss下降很快,但后期fine-tune时模型在动作分布上一直有奇怪的微小抖动,最终评测成功率反而低了3个百分点。
后来换成峰值学习率1e-4,scheduler调到慢热长衰减,效果反而更好。这说明模仿学习任务里,学习率调优不是单纯追求下降速度快,更重要的是让模型在收敛阶段有足够长的时间稳定到最优动作分布上。scheduler的价值这时候就体现出来了:控制模型在整个训练周期里什么阶段该快、什么阶段该慢。
6. 写在最后的调参体会
LeRobot里scheduler的实现本身不算复杂,核心也就是一个get_scheduler加训练循环里的step()调用。但把所有细节串起来以后会发现,真正决定训练效果的恰恰是这些"不复杂"的部分:warmup步数占总步数比例、梯度累积时的更新计数、checkpoint恢复时的状态对齐、分布式场景下加速器接管。我的建议是不要盲目修改scheduler参数,先把默认配置跑完整一遍,记录学习率曲线,再根据曲线去判断哪个阶段需要调,一次只动一个参数。
我自己的习惯是每开一个新任务,先跑50步快速验证scheduler曲线形状对不对,再决定要不要继续。scheduler这个事情,看起来小,但它每天都在替你做决定,模型学得快不快、稳不稳,大半的功劳和锅都跟它有关。希望这篇拆解能帮你少走一些弯路。