
什么是 FlashREINFORCE?
FlashREINFORCE 是一个无需评论家模型、用于异步训练智能体语言模型的强化学习框架。它专为长时域工作负载而设计,在这类工作负载中,由于推理长度、工具调用或环境交互存在差异,轨迹会在不规则的时间完成。
FlashREINFORCE 不会为每个提示生成多个同源 rollout,而是对每个提示使用一次 rollout。因此,一个包含 B 条轨迹的批次可以覆盖 B 个不同的提示,已完成的轨迹也无需等待同一提示的其他响应,就能传递给学习器。
该仓库提供独立的 PyTorch 参考损失、CPU 优化器更新示例、测试、用于 R1 和 Qwen2.5-Math 的 Molt 启动器,以及针对推理、Python 工具、混合专家模型和 ALFWorld 的实验配置。完整的异步训练启动器使用 Molt。
强化学习就应该做 REINFORCE。
为什么使用单次 rollout?
组内相对方法依赖于同一提示生成的多个 rollout。在异步智能体工作负载中,这一要求可能引入同步屏障,并减少固定 rollout 预算所能覆盖的不同提示数量。
单次 rollout 学习移除了同源 rollout 的要求,但也带来了三个稳定性问题:
- 没有同源轨迹组可用于构造提示级基线。
- 异步收集的轨迹可能由较旧的行为策略生成。
- 如果每个 token 都获得相同的批次级权重,较长的失败轨迹可能主导一次更新。
FlashREINFORCE 通过 One-Batch REINFORCE、Sequence Trust Region 和 Sample-Mean Optimization 解决这些问题。
核心组件
1. One-Batch REINFORCE
对于每个新批次,学习器会在相互独立的提示之间对标量奖励进行中心化。若 R_i 是轨迹 i 的奖励,则其优势为该轨迹奖励减去批次均值:
mean_reward = sum(R_j for j in batch) / B
A_i = R_i - mean_reward
高于均值的轨迹获得正优势,低于均值的轨迹获得负优势。这样无需训练评论家模型即可提供带符号的反馈。学习器在该批次上执行一次优化器更新,然后丢弃该批次。
2. Sequence Trust Region
每条轨迹都必须保留实际生成每个采样 token 的行为策略概率。学习器重新计算当前策略概率,并构造 token 重要性比率:
rho_i_t = exp(current_log_probability_i_t
- behavior_log_probability_i_t)
该比率会校正存储历史状态下的采样动作分布。但它无法校正历史本身来自较旧行为策略这一事实,因此 FlashREINFORCE 还会评估完整轨迹上累积的策略漂移。
对于每个采样动作,它根据行为概率 p 和当前概率 q 计算 Bernoulli KL 代理值:
d_i_t = p * log(p / q)
+ (1 - p) * log((1 - p) / (1 - q))
token 值会在轨迹上取平均。只有当其平均代理值不超过配置的信任阈值时,轨迹才会被接纳:
mean_drift_i = sum(d_i_t for t in trajectory_i) / T_i
mask_i = 1 if mean_drift_i <= delta else 0
被拒绝的轨迹不会贡献任何 token 损失。因此,信任决策作用于完整的已存储历史,而不是彼此孤立的 token。
3. Sample-Mean Optimization
token 级批次平均会自动让较长轨迹产生更大影响。当一条较长的失败 rollout 包含有用的中间步骤、但最终以负奖励结束时,这可能造成不利影响。
FlashREINFORCE 首先对每条轨迹内的 token 贡献取平均,然后再对轨迹取平均:
trajectory_term_i = (
mask_i * advantage_i / trajectory_length_i
* sum(importance_ratio_i_t for each token t)
)
objective = sum(trajectory_term_i for i in batch) / B
因此,无论响应长度如何,每条轨迹获得的权重都是 1/B。实际损失使用已分离梯度的 importance ratio。基础方法不使用学习到的评论家模型、比率裁剪或参考模型前向传播。
安装和本地设置
克隆或下载 FlashREINFORCE 仓库,进入其根目录,并连同测试依赖以可编辑模式安装:
pip install -e '.[test]'
接下来运行随附的 CPU 示例:
python examples/toy_update.py
该示例会检查参考目标是否能够在 CPU 上执行一次优化器更新。它是功能验证示例,并非论文基准结果的复现。
最后运行测试套件:
python -m pytest -q
示例和测试均成功运行,说明本地参考实现及其依赖可用。
探索参考实现
核心参考目标位于 flashreinforce/loss.py。它涵盖论文和 README 中描述的操作:
- 跨批次的奖励中心化。
- token 级重要性采样。
- 序列级信任接纳。
- 样本均值优化。
- 基于可选熵的失败 token 过滤。
先从 examples/toy_update.py 开始,了解参考损失如何参与优化器更新。然后查看 flashreinforce/loss.py,将示例张量和掩码与目标的各个阶段对应起来。
CPU 示例是最安全的基础使用路径,因为它不需要完整的分布式部署基础设施。
了解一次异步更新
生产环境中的 FlashREINFORCE 更新包含六个阶段:
- 收集:部署工作进程彼此独立地为每个提示提交一条已完成的轨迹。每条轨迹都包含其奖励,以及用于采样其中各个 token 的真实行为策略概率。
- 批处理:学习器获取接下来的 B 条已完成轨迹,并将其奖励进行中心化,以得到带符号的优势值。
- 校正:学习器重新计算当前策略的对数概率,并构造 token 重要性比率。
- 信任筛选:它计算每条轨迹的采样动作平均 KL 代理值,只有当该值处于配置的阈值范围内时,才接纳完整轨迹。
- 优化:它先在每条获接纳的轨迹内部求平均,再在整个批次上求平均,随后执行一次优化器更新。
- 丢弃:它不会使用已收集的批次执行额外的学习器更新。
最后的丢弃步骤是该方法的重要组成部分。FlashREINFORCE 对每个新收集的批次执行一次完整批次更新,而不是基于同一个行为策略快照运行多次连续的小批次更新。
预览 Molt 训练配置
仓库包含 scripts/train_molt.py,用于通过 Molt 的异步实现启动训练。在分配完整训练任务之前,请使用其试运行模式预览训练参数:
python scripts/train_molt.py --dry-run
有关 GPU 设置、数据准备、训练命令、论文概念到代码的映射,以及已知的复现限制,请参阅 docs/training.md。examples/README.md 指南包含任务预设、消融实验,以及带功能检查的原生训练器启动器。
R1 和 Qwen2.5-Math 启动器已固定在仓库中。工具使用和 ALFWorld 设置还需要兼容的智能体与环境,因此在启动这些实验前,请先查看其运行时要求。
每条轨迹必须保留的数据
该更新依赖于生成过程中捕获的信息。至少,概念上的训练流程需要保留已完成的 token 序列、其标量奖励,以及每个采样动作对应的行为概率。随后,学习器通过评估已存储的动作及其历史记录,获取对应的当前策略概率。
不要使用事后重建的概率替换已记录的采样概率。后续的模型评估不一定能复现推理引擎在采样时实际使用的概率。
高级稳定性建议
保留采样支持集
精确的重要性校正假设目标策略不会在行为策略支持集之外分配概率质量。激进的 top-k 或 top-p 截断可能违反这一假设。因此,应结合支持集要求选择采样设置。
不要将 token 重要性采样视为完整轨迹校正
Token 重要性比率会校正以已存储历史为条件的采样动作,但不会校正这些历史的分布。请使用序列信任区域筛除行为策略与学习器之间累计偏移过大的轨迹。
保持更新新鲜
每收集一个批次就执行一次优化器更新,然后将其丢弃。对同一个过时批次执行多次学习器更新,偏离了文档所述的一次性方法,并可能加剧策略不匹配。
保留轨迹级加权
先在每条轨迹内部对 token 贡献求平均,再在整个批次中对各条轨迹求平均。将所有有效 token 展平后计算整个批次的平均值,会重新引入响应长度乘数,使较长轨迹获得不成比例的权重。
正确理解信任掩码
序列信任区域会针对整条轨迹作出一次接纳决策。如果某条轨迹超过阈值,则该更新会移除其中所有 token 的损失。
区分验证与复现
CPU 示例只能验证一次优化器更新。要复现研究结果,需要按照训练指南和示例指南中记录的 GPU、数据、模型、智能体、环境及异步运行时配置进行设置。
报告的应用领域
报告中的实验涵盖数学推理、多轮 Python 工具使用、混合专家训练,以及交互式 ALFWorld 决策。README 报告了 DeepSeek-R1-Distill-Qwen-1.5B 在 6,000 次更新中稳定进行长链式思维训练的结果、Qwen2.5-Math-1.5B 的数学推理结果、使用 Qwen2.5-7B-Instruct 和 Qwen3-30B-A3B 的 Python 工具实验,以及使用 Qwen2.5-7B-Instruct 的 ALFWorld 实验。
这些结果展示了该框架预期的适用范围,但不应期待通过小型 CPU 更新示例获得相同结果。
实用资源
结语
FlashREINFORCE 将单次 rollout 的提示词覆盖与紧凑的无 critic 更新相结合。以批次为中心的奖励提供带符号的优势值,token 重要性采样用于处理过时动作,Sequence Trust Region 控制累积的轨迹漂移,而样本均值优化则避免响应长度决定轨迹权重。
请从可编辑安装、CPU 更新示例和测试开始。确认这些内容正常运行后,使用 --dry-run 预览 Molt 启动器,并按照训练指南完成特定于模型的 GPU 与数据准备。项目采用 Apache License 2.0 授权。
