Discrete Flow Matching for Offline-to-Online Reinforcement Learning

📄 arXiv: 2605.12379v1 📥 PDF

作者: Fairoz Nower Khan, Nabuat Zaman Nahim, Peizhong Ju

分类: cs.LG, cs.AI

发布日期: 2026-05-12


💡 一句话要点

DRIFT:用于离线到在线强化学习的离散流匹配方法

🎯 匹配领域: 支柱二:RL算法与架构 (RL & Architecture)

关键词: 离散动作空间 强化学习 离线到在线学习 流匹配 连续时间马尔可夫链

📋 核心要点

  1. 现有生成策略方法主要针对连续控制,难以直接应用于离散动作空间的强化学习任务。
  2. DRIFT通过优势加权的离散流匹配损失,在线微调离线预训练的连续时间马尔可夫链策略。
  3. 实验表明,DRIFT在离散动作RL任务中实现了稳定的离线到在线改进,并在Jericho上取得了最佳平均分数。

📝 摘要(中文)

许多强化学习(RL)任务具有离散动作空间,但大多数基于扩散和流匹配的生成策略方法是为连续控制设计的。同时,生成策略通常严重依赖离线数据集,而离线到在线RL本身就具有挑战性,因为策略必须通过新的交互来改进,同时不丢失从静态数据中学到的有用行为。为了解决这些挑战,我们引入了DRIFT,这是一种在线微调方法,它使用优势加权的离散流匹配损失来更新离线预训练的连续时间马尔可夫链(CTMC)策略。为了保留有用的预训练知识,我们添加了一个路径空间惩罚,它正则化了完整的CTMC轨迹分布,而不仅仅是最终的动作分布。对于大型离散动作空间,我们引入了一种候选集近似,它通过从参考策略rollout和均匀探索中采样的一小部分动作来更新actor。我们的理论分析表明,候选集误差受缺失目标概率质量的控制,并且随着候选集覆盖更多高概率动作,诱导的CTMC生成器误差会减小。在流行的离散动作RL任务上的实验表明,我们的方法在所有任务中都提供了稳定的离线到在线改进,在使用简单的GRU编码器的情况下,在Jericho上实现了最高的平均分数,同时优于使用预训练语言模型的方法。受控实验进一步证实,路径空间惩罚在微调期间保持有界,并且CTMC生成器比确定性基线更快地适应转移的奖励。候选集机制得到了稳定性分析的支持,该分析表明生成器误差随着候选覆盖率呈指数下降。

🔬 方法详解

问题定义:论文旨在解决离散动作空间下的离线到在线强化学习问题。现有方法,特别是基于扩散模型或流匹配的生成策略,主要针对连续动作空间设计,难以直接应用于离散动作空间。此外,如何在利用离线数据学习到的知识的同时,通过在线交互进行策略改进,是一个挑战。

核心思路:论文的核心思路是利用连续时间马尔可夫链(CTMC)来表示策略,并使用离散流匹配损失进行在线微调。通过优势加权的方式,鼓励策略向更有利的方向更新。为了防止在线微调过程中遗忘离线数据学习到的知识,引入了路径空间惩罚,对整个CTMC轨迹分布进行正则化。

技术框架:DRIFT方法包含以下主要步骤:1) 使用离线数据集预训练一个CTMC策略。2) 在线微调阶段,使用优势加权的离散流匹配损失更新策略。3) 为了保留离线知识,添加路径空间惩罚项。4) 对于大型离散动作空间,采用候选集近似方法,只更新从参考策略rollout和均匀探索中采样的一小部分动作。

关键创新:论文的关键创新在于:1) 将连续时间马尔可夫链应用于离散动作空间的强化学习。2) 提出了优势加权的离散流匹配损失,用于在线微调。3) 引入了路径空间惩罚,用于保留离线学习到的知识。4) 提出了候选集近似方法,用于处理大型离散动作空间。

关键设计:路径空间惩罚通过正则化整个CTMC轨迹分布来保留离线知识。候选集近似通过从参考策略rollout和均匀探索中采样动作来减小计算复杂度。优势加权确保策略向更有利的方向更新。损失函数结合了离散流匹配损失和路径空间惩罚项。网络结构方面,可以使用GRU等循环神经网络作为编码器。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明,DRIFT方法在多个离散动作RL任务中取得了显著的性能提升,尤其是在Jericho游戏中,使用简单的GRU编码器就超越了使用预训练语言模型的方法,实现了最高的平均分数。受控实验还验证了路径空间惩罚的有效性,以及CTMC生成器对奖励变化的快速适应能力。候选集机制的稳定性分析表明,生成器误差随着候选覆盖率呈指数下降。

🎯 应用场景

该研究成果可应用于各种具有离散动作空间的强化学习任务,例如游戏AI、机器人控制、推荐系统、对话系统等。通过离线数据预训练和在线微调,可以快速提升策略性能,并适应环境变化。该方法在需要从历史数据中学习经验,并不断与环境交互进行优化的场景下具有重要价值。

📄 摘要(原文)

Many reinforcement learning (RL) tasks have discrete action spaces, but most generative policy methods based on diffusion and flow matching are designed for continuous control. Meanwhile, generative policies usually rely heavily on offline datasets and offline-to-online RL is itself challenging, as the policy must improve from new interaction without losing useful behavior learned from static data. To address those challenges, we introduce DRIFT, an online fine-tuning method that updates an offline pretrained continuous-time Markov chain (CTMC) policy with an advantage-weighted discrete flow matching loss. To preserve useful pretrained knowledge, we add a path-space penalty that regularizes the full CTMC trajectory distribution, rather than only the final action distribution. For large discrete action spaces, we introduce a candidate-set approximation that updates the actor over a small subset of actions sampled from reference-policy rollouts and uniform exploration. Our theoretical analysis shows that the candidate-set error is controlled by missing target probability mass, and the induced CTMC generator error decreases as the candidate set covers more high-probability actions. Experiments on prevailing discrete action RL task show that our method provides stable offline-to-online improvement across all tasks, achieving the highest average score on Jericho with a simple GRU encoder while outperforming methods that use pretrained language models. Controlled experiments further confirm that the path-space penalty remains bounded during fine-tuning and that the CTMC generator adapts to shifted rewards faster than deterministic baselines. The candidate-set mechanism is supported by a stability analysis showing that the generator error decreases exponentially with candidate coverage.