Diffusion Policy训练手记（上篇）：聚焦训练流程——从Batch梯度更新到采样策略的解惑之路
作为一名深耕机器人控制领域的开发者，最近在基于 Diffusion Policy 开发工业流水线机械臂高精度控制方案时，我被训练流程里的几个核心问题卡住了。这些问题看似基础，却直接决定了模型训练的效率和方向。今天把我从“一头雾水”到“豁然开朗”的学习过程整理出来，希望能帮到同样在Diffusion Policy训练中摸爬滚打的同学。
我的任务目标很明确：训练模型根据过去2帧图像+当前状态，预测未来16步的专家动作，实现产线工件的高精度抓取。设定 batch_size=64 ，但训练前的几个疑问，让我迟迟不敢启动训练脚本。
一、 第一个核心困惑：64个样本的梯度，到底怎么更新参数？
最开始我陷入了一个误区：一个batch里有64个独立样本，每个样本都会计算出一套梯度，模型该用哪一套梯度更新参数？难不成64套梯度会基于同一个初始参数“各自为政”，最后导致参数更新混乱？
带着这个疑问，我翻遍了Mini-Batch梯度下降的核心原理，才彻底理清了逻辑——64个样本的梯度不会各自更新，而是先求平均，再完成一次参数更新。
具体的执行步骤可以拆成三步：
1. 把包含64个样本的batch喂入模型，针对每个样本，计算“预测噪声”与“真实噪声”的MSE损失，再通过反向传播得到64组独立的参数梯度；
2. 对这64组梯度，在每个参数维度上分别计算平均值，得到一个全局平均梯度；
3. 结合学习率，用这个全局平均梯度对模型参数完成一次更新。
用公式可以更直观地表达这个过程（以单个参数 w 为例）：
\Delta w = -\eta \cdot \frac{1}{64}\sum_{i=1}^{64}\nabla_{w}L_i
其中 \eta 代表学习率，L_i 是第 i 个样本的损失值，\nabla_{w}L_i 是第 i 个样本对参数 w 的梯度。
这个逻辑彻底打消了我的顾虑：一个batch只对应一次梯度下降，既不会出现参数更新“打架”的情况，还能兼顾训练效率与梯度稳定性——相比单样本随机梯度下降（SGD），它的梯度更平滑，训练波动更小；相比全量梯度下降，它更节省显存，训练速度更快。
二、 第二个核心困惑：Batch该从单个episode采样，还是全局随机采样？
解决了梯度更新的问题，新的疑问又来了：我的数据集包含1000个产线作业的完整episode，每个episode里有若干组“2帧图像+状态+16步动作”的样本。采样时，是把一个batch的64个样本限定在单个episode里，还是从所有episode里随机抽取？
这个问题的答案，直接关系到数据多样性，进而影响模型的学习效果。我通过对比两种采样方式的优劣，找到了最优解：必须选择全局随机采样。
如果限定在单个episode内采样，模型很容易“死记硬背”这个episode的特殊动作细节——比如某次抓取中机械臂的偶然轨迹、某个特定光照下的工件特征，导致模型学不到跨场景的通用控制规律。
而全局随机采样的逻辑是，先把1000个episode的所有样本汇总成一个大列表，再从中随机抽取64个样本组成batch。这样做能让每个batch都覆盖不同的工件位置、光照强度、机械臂初始姿态，迫使模型聚焦于“机械臂如何根据观测调整动作”的核心逻辑，而非单个episode的“专属剧本”。
这里附上我实际使用的采样伪代码，方便大家直接参考：
python
import random
def sample_batch(dataset, batch_size=64):
"""
从所有episode中全局随机采样一个batch
dataset格式：{
"episode_0": [样本0, 样本1, ...],
"episode_1": [样本0, 样本1, ...],
...
"episode_999": [样本0, 样本1, ...]
}
"""
all_samples = []
# 汇总所有episode的样本
for episode_samples in dataset.values():
all_samples.extend(episode_samples)
# 全局随机选取batch_size个样本
batch_samples = random.sample(all_samples, batch_size)
return batch_samples
三、 第三个核心困惑：训练时需要跑多步去噪吗？
在理清采样和梯度更新的逻辑后，我又对Diffusion Policy的训练核心产生了疑问：Diffusion Policy的本质是逐步去噪还原动作，那训练时是不是也要像推理时一样，跑满所有时间步？比如我设定了100个扩散时间步，训练时是不是要循环100次去噪？
这个疑问的答案，直接点透了Diffusion Policy 训练与推理的核心区别：训练时只需要预测单步噪声，不需要跑多步；只有推理时才需要逐步去噪。
具体来说，训练的核心逻辑是“随机选步、单步预测”：
1. 对batch里的每个样本，随机分配一个扩散时间步 k（比如 k=35 或 k=88）；
2. 模型输入“第 k 步的加噪动作 + 当前观测 + 时间步 k”，只需要预测第 k 步这一层的噪声；
3. 用预测噪声和真实噪声的MSE损失更新模型，让模型学会“给定任意时间步，都能精准还原这一步的噪声”的能力。
之所以不需要训练时跑多步，是因为模型的目标是掌握“单步去噪”的通用能力，而不是记住某一条去噪路径。当模型在成千上万的随机时间步上完成训练后，推理时就可以从全噪声状态出发，逐步去噪，最终还原出高精度的动作序列。
四、 训练流程梳理：从采样到参数更新的完整链路
理清了所有核心疑问后，我把Diffusion Policy的训练流程梳理成了一套可直接落地的步骤：
1. 数据集准备：汇总所有episode的样本，每个样本包含“2帧图像+状态+16步专家动作”；
2. Batch采样：调用上述伪代码，从全局样本池中随机抽取64个样本组成batch；
3. 随机分配时间步：给batch内每个样本分配随机时间步 k，生成对应加噪动作；
4. 前向计算与损失计算：模型输入“加噪动作+观测+时间步 k”，预测单步噪声，计算MSE损失；
5. 梯度平均与参数更新：计算64个样本的梯度平均值，结合学习率完成一次参数更新；
6. 重复迭代：循环上述步骤，直到训练集loss趋于稳定。
写在最后
这篇手记聚焦的是Diffusion Policy训练流程的核心疑问，解决的是“怎么训”的问题。下一篇，我会继续分享后续遇到的过拟合难题，以及由此提炼出的模型、数据、正则化三者的底层逻辑。
