Paper Reading: Offline Reinforcement Learning for LLM Multi-Step Reasoning

Wang, Huaijie, Shibo Hao, Hanze Dong, Shenao Zhang, Yilin Bao, Ziran Yang和Yi Wu. 《Offline Reinforcement Learning for LLM Multi-Step Reasoning》. arXiv, 2024.12

Code released:OREO Huggingface Model: Policy|Value

Motivation:

DPO的局限性

需要成对数据：DPO需要成对的偏好数据，而多步推理任务通常只有稀疏奖励的独立轨迹。
缺乏信用分配：DPO在整个轨迹上进行优化，无法对关键token进行精细的信用分配。（从此点入手）

传统的强化学习方法（如DPO）在整个轨迹上进行优化，无法区分哪些步骤对最终结果贡献更大，导致信用分配不够精细。

OREO能否进行精细的信用分配呢？请参考下文详细解释

OREO的动机

基于上述分析，作者提出了OREO方法，旨在解决DPO的局限性：

无需成对数据：OREO可以直接利用稀疏奖励的离线数据集。
更好的信用分配：通过联合学习策略和价值函数，OREO能够更精细地分配奖励，特别适合多步推理任务。

Method:

OREO方法简介

基于SAC的方法（增加探索度）

OREO是一种基于最大熵强化学习（RL）原理的离线强化学习方法。它通过优化软贝尔曼方程（soft Bellman equation）来同时学习策略模型和价值函数。

这种方法利用预先收集的数据进行训练，无需与环境进行实时交互(真实世界训练成本高)。最大熵RL框架鼓励探索，通过最大化策略的熵来实现更健壮的策略。在OREO中，软贝尔曼方程可能包含熵正则化项，这在标准贝尔曼方程的基础上增加了策略熵的考虑，以平衡探索与利用。

这种联合学习策略和价值函数的方式，有助于在多步推理任务中实现更好的信用分配和性能提升。OREO相较于其他方法如DPO，在处理复杂任务时表现出更强的适应性和有效性。

需要理解的问题

3 Preliminaries

将LLM推理任务形式化为MDP。

引入KL正则化的强化学习目标。

基于软贝尔曼方程的最大熵强化学习框架。

1.OREO的核心思想

OREO是一种离线强化学习（Offline RL）方法，旨在通过优化软贝尔曼方程（Soft Bellman Equation）来提升大语言模型（LLMs）在多步推理任务中的表现。OREO的核心思想包括：

联合学习策略模型和价值函数：通过同时优化策略模型 $\pi_\theta$ 和价值函数 $V_\phi$，OREO能够更好地进行信用分配（Credit Assignment）。
无需成对数据：与DPO不同，OREO可以直接利用稀疏奖励的离线数据集，而无需成对的偏好数据。
测试时搜索：训练得到的价值函数可以在推理时引导搜索过程（如束搜索），进一步提升推理性能。

2.学习目标（Learning Objective）

OREO的学习目标是通过优化软贝尔曼方程来联合学习策略模型和价值函数。具体目标如下：

（1）价值函数的训练

价值函数 $V_\phi$ 的训练目标是预测从当前状态 $s_t$ 开始的未来累积奖励。其损失函数为均方误差（MSE）：

\[L_V(\phi) = \frac{1}{T} \sum_{t=0}^{T-1} \left( V_\phi(s_t) - R_t + \beta \sum_{i \geq t} \log \frac{\pi_\theta(a_i | s_i)}{\pi_{\text{ref}}(a_i | s_i)} \right)^2\]

其中：

$V_\phi(s_t)$ 是价值函数对状态 $s_t$ 的预测值。
$R_t$ 是从时间步 $t$ 开始的真实累积奖励。
$\pi_\theta$ 是当前策略模型，$\pi_{\text{ref}}$ 是参考策略。
$\beta$ 是KL正则化项的权重。

（2）策略模型的训练

策略模型 $\pi_\theta$ 的训练目标是通过优化软贝尔曼方程来改进策略。其损失函数为：

\[L_\pi(\theta) = \frac{1}{T} \sum_{t=0}^{T-1} \left( V_\phi(s_t) - R_t + \beta \log \frac{\pi_\theta(a_t | s_t)}{\pi_{\text{ref}}(a_t | s_t)} + \text{sg}\left[\beta \sum_{i > t} \log \frac{\pi_\theta(a_i | s_i)}{\pi_{\text{ref}}(a_i | s_i)}\right] \right)^2 + \alpha L_{\text{reg}}\]

其中：

$\text{sg}[\cdot]$ 表示停止梯度操作，确保梯度不会传播到后续时间步。
$L_{\text{reg}}$ 是KL正则化项，用于防止策略模型过度偏离参考策略：

\[L_{\text{reg}} = \frac{1}{T} \sum_{t=0}^{T-1} \text{KL}\left[\pi_\theta(\cdot | s_t) \parallel \pi_{\text{ref}}(\cdot | s_t)\right]\]

$\alpha$ 是正则化项的权重。

3. 损失函数变体（Loss Variants）

为了适应不同的任务需求，OREO提供了两种损失函数变体：

（1）步级OREO（Step-level OREO）

在步级OREO中，一个动作被视为一个完整的推理步骤（而不是单个token）。其损失函数与标准OREO类似，但动作的定义和概率计算方式有所不同。

（2）响应级OREO（Response-level OREO）

响应级OREO类似于DPO，只在初始状态上进行优化。其损失函数为：

\[L_\pi^{\text{resp}}(\theta) = \left( V_\phi(s_0) - R_0 + \beta \sum_{i \geq 0} \log \frac{\pi_\theta(a_i | s_i)}{\pi_{\text{ref}}(a_i | s_i)} \right)^2 + \alpha L_{\text{reg}}\]

4. 迭代OREO（Iterative OREO）

OREO可以扩展为迭代训练框架，具体步骤如下：

初始训练：使用初始离线数据集训练策略模型和价值函数。
数据生成：使用训练后的策略模型生成新的轨迹数据。
数据筛选：选择高质量的轨迹（如成功轨迹）加入训练集。
迭代训练：重复上述步骤，逐步提升模型性能。

5. 测试时搜索（Test-Time Search）

在推理阶段，OREO利用训练得到的价值函数 $V_\phi$ 来引导搜索过程，进一步提升推理性能：

数学推理任务：使用束搜索（Beam Search），在每一步选择价值函数预测值最高的候选动作。
Embodied Agent任务：使用Best-of-K策略，在每一步采样 $K$ 个动作，选择价值函数预测值最高的动作。

OREO是如何进行精细信用分配的

OREO通过以下方式实现更精细的信用分配：

（1）联合学习策略模型和价值函数

OREO同时训练策略模型 $πθ*πθ*$ 和价值函数 $Vϕ*Vϕ*$：

策略模型：负责生成动作（即推理步骤中的token）。
价值函数：评估当前状态的未来累积奖励。

通过联合学习，OREO能够在每个时间步上评估当前状态的价值，并根据价值函数的反馈调整策略模型，从而更精确地分配奖励。

（2）软贝尔曼方程的优化

OREO通过优化软贝尔曼方程来实现信用分配：

\[V^*(s_t) - V^*(s_{t+1}) = r(s_t, a_t) - \beta \log \frac{\pi^*(a_t | s_t)}{\pi_{\text{ref}}(a_t | s_t)}\]

其中：

$V^*(s_t)$ 和 $V^*(s_{t+1})$ 分别是当前状态和下一个状态的最优价值。
$r(s_t, a_t)$ 是当前状态-动作对的奖励。
$\beta \log \frac{\pi^*(a_t | s_t)}{\pi_{\text{ref}}(a_t | s_t)}$ 是KL正则化项，用于控制策略与参考策略的偏离。

通过优化软贝尔曼方程，OREO能够在每个时间步上评估动作的贡献，并将奖励分配到对最终结果影响最大的关键步骤。

（3）时间步级别的优化

与DPO在整个轨迹上进行优化不同，OREO在每个时间步上进行优化。这使得OREO能够更精细地分析每个动作对最终结果的贡献，从而更好地进行信用分配。