Reinforcement Learning for Large Language Models: An Overview

From reinforcement learning foundations to policy optimization for LLMs.

Author

Updated

Jul, 31, 2026

Category

Introduction

本文围绕一个核心问题展开:

我们是如何使用强化学习来提高大语言模型的能力上限的?

为了解决这个问题,我们需要回答:

  1. 强化学习是如何建模的?
  2. 强化学习有哪些算法?哪些强化学习算法适用/不适用于大语言模型?这些算法背后的原理是什么?
  3. 现代大语言模型使用的强化学习算法是什么?

我们将按照下面的学习路径组织:

  1. Foundations: 我们将从基本的 定义 (MDP, return) 出发,建立 LLM generation 与 RL 的对应关系,再用 Bellman Equation 将长期回报转化为递推问题。
  2. Evaluation and Improvement: 这一部分,我们先回答“当前 policy 有多好” (policy evaluation), 再介绍 value-based methods, 该部分会解释为什么 Q-learning/DQN 对理解 RL 很重要,但通常不是优化 LLM policy 的最终选择。
  3. Policy Optimization: 这一部分,我们从 policy gradient methods 逐步过渡到 actor-critic methods, 在这个过程中,我们将逐步优化 policy model 和 value model 的训练与使用。这也是后续 PPO 算法的核心。
  4. Modern RL Algorithms for LLM: 我们将从 TRPO 开始, 逐步介绍 PPO, GRPO 所做的改进,到这一步,我们就接触到了现代 RL 算法的核心部分。

前两部分建立后续推导所需的共同语言;如果读者已经熟悉经典 RL,可以先阅读 LLM as an MDP,再从 Policy Evaluation 进入算法主线。

Course materials

Acknowledgements

这个 tutorial 主要参考了以下两位老师的课程和材料:

RL Basics

本章的目标是建立全文统一使用的语言:我们先用 MDP 描述 agent 与 environment 的交互,再定义 trajectory, return 和优化目标。 最后我们把 auto-regressive generation 写成一个 token-level MDP. 后续章节我们会一直使用这里定义的 LLM state, action 和 reward.

RL 的基本思想是让 agent 通过与环境交互,学习能够最大化期望回报的策略。 其定义为:

强化学习是一个通过构建可以与 environment 进行交互的 agent 来解决控制和决策任务的学习框架,交互的方式为 agent 执行 action, 然后 environment 给予反馈。

RL 的执行过程如下所示

The RL process (source from Sutton)

RL 的数学建模依赖于 Markov Decision Process, 下面我们先介绍相关概念。

Markov Decision Process (MDP)

Definition: Markov Decision Process

MDP 是一个可以用于描述强化学习的数学模型,形式化为:

  1. 状态 (State): stSs_t \in \mathcal{S}, S\mathcal{S} 为状态空间。
  2. 动作 (Action): atAa_t \in \mathcal{A}, A\mathcal{A} 为动作空间。
  3. 奖励 (Reward): rtRr_t \in \mathbb{R}, 环境对 agent 执行动作 ata_t 的反馈。
  4. 终止时间 (Terminal Time): TT,且定义 sT=terms_T = \langle\text{term}\rangle 为终止状态。
  5. 初始状态分布 (Initial State): s0p0s_0 \sim p_0.
  6. 状态转移概率 (Transition Probability): (rt,st+1)p(,st,at)(r_t, s_{t+1}) \sim p(\cdot, \cdot \mid s_t, a_t).
  7. 马尔可夫性 (Markov Property):
p(st+1st,at,)=p(st+1st,at), p(rtst,at,)=p(rtst,at).\htmlId{markov_property}{\begin{equation} p(s_{t+1} \mid s_t, a_t, \ldots) = p(s_{t+1} \mid s_t, a_t),\ p(r_t \mid s_t, a_t, \ldots) = p(r_t \mid s_t, a_t). \end{equation}}

为了简化,我们采用以下约定:

  1. rtr_t 完全由 (st,at)(s_t,a_t) 决定时,记为 rt=r(st,at)r_t = r(s_t, a_t).
  2. 状态转移函数是平稳的 (stationary), 即 pt(r,ss,a)=p(r,ss,a)p_t(r,s'\mid s,a) = p(r,s'\mid s,a).
  3. ata_t 由策略 π\pi 决定,确定性策略记为 at=π(st)a_t = \pi(s_t), 随机策略记为 atπ(st)a_t \sim \pi(\cdot \mid s_t).
  4. 通常我们使用神经网络表示策略,即 π=πθ\pi = \pi_\theta, 其中 θ\theta 为网络参数。
  5. 我们用 StS_t, AtA_t 代表 tt 时刻的状态和动作,这是一个随机变量,用 St=stS_t=s_t, At=atA_t=a_t 代表当前状态为 sts_t, 当前动作为 ata_t.

Trajectory

接下来,我们介绍 trajectory, 也就是 RL 算法所需要的数据的定义

Definition: Trajectory

Trajectory (rollout) 定义为 agent 与环境完整的一次交互过程:

τ=(s0,a0,r0,s1,a1,r1,,sT1,aT1,rT1,sT)\tau = (s_0, a_0, r_0, s_1, a_1, r_1, \ldots, s_{T-1}, a_{T-1}, r_{T-1}, s_T)

其概率分布为

p(τs0,π,p)=p0(s0)t=0T1π(atst)p(rt,st+1st,at)p(\tau \mid s_0, \pi, p) = p_0(s_0) \prod_{t=0}^{T-1} \pi(a_t \mid s_t) \, p(r_t, s_{t+1} \mid s_t, a_t)
  • 轨迹的概率分布由初始状态分布、策略以及环境模型共同决定,记为 τ(p0,π,p)\tau \sim (p_0, \pi, p).
  • 轨迹的概率分布使用了 Markov property.
  • 本教程仅考虑 finite horizon MDP, 即 T<T < \infty.

Return

接下来,我们介绍 RL 算法的优化目标 (discounted) return.

Definition: Discounted Return

一条 trajectory 上的 (discounted) return 定义为:

R(τ)=t=0T1γtrtR(\tau) = \sum_{t=0}^{T-1} \gamma^t r_t

其中 γ(0,1]\gamma \in (0, 1]discount factor,用于将未来的奖励折现到当前时刻。

从时刻 tt 开始的 return 记为:

Gt=k=tT1γktrkG_t = \sum_{k=t}^{T-1} \gamma^{k-t} r_k

GtG_t 存在如下递推关系:

Gt=rt+γGt+1G_t = r_t + \gamma G_{t+1}

易知

R(τ)=G0.R(\tau) = G_0.

Objective of RL

RL 的最终目标为最大化 expected return, 形式化为:

maxπEτ(p0,π,p)R(τ)=maxπEτ(p0,π,p)[t=0T1γtrt]\htmlId{rl_objective}{\begin{equation} \max_{\pi} \quad \mathbb{E}_{\tau \sim (p_0, \pi, p)} \, R(\tau) = \max_{\pi} \quad \mathbb{E}_{\tau \sim (p_0, \pi, p)} \left[ \sum_{t=0}^{T-1} \gamma^t r_t \right] \end{equation}}

这里期望 Eπ\mathbb{E}^\pi 对应的随机变量为 atπ(st)a_t \sim \pi(\cdot \mid s_t), (rt,st+1)p(,st,at)(r_t, s_{t+1}) \sim p(\cdot, \cdot \mid s_t, a_t).

这个目标函数源于强化学习的核心假设

Reward Hypothesis: 所有的目标都可以被描述为最大化 expected return.

由于实际问题中,初始状态 s0s_0 和状态转移函数 pp 一般保持不变,因此我们一般省略目标函数中的 s0s_0pp, 将目标函数简化为

maxπEτπR(τ)\max_{\pi} \quad \mathbb{E}_{\tau\sim\pi} \, R(\tau)

Connection with LLM

这一节我们介绍如何用 RL 的定义来建模 auto-regressive LLM. LLM 定义如下:

ytLLM(y<t,x)y_t \sim \mathrm{LLM}(\cdot \mid y_{<t}, x)

其中 xx 是输入的 prompt, y<ty_{<t} 是已经生成的部分 (generated prefix), yty_t 是下一个要生成的 token.

RL 和 LLM 的对应关系如下 (Shen et al., 2026)

RLLLMDescription
initial state s0s_0xxprompt
state sts_tst=(y<t,x)s_t=(y_{<t}, x)prompt and generated prefix
action ata_tyty_tnext token selected from vocabulary
policy π(atst)\pi(a_t\mid s_t)LLM(ytst)\mathrm{LLM}(y_t \mid s_t)next token prediction distribution
trajectory τ\tau(y1:T,x)(y_{1:T},x)complete prompt-response pair
Return R(τ)R(\tau)R(x,y)R(x,y)reward model or verifier score
transition P(st+1st,at)P(s_{t+1}\mid s_t,a_t)st+1=concat(yt,st)s_{t+1} = \mathrm{concat}(y_t, s_t)deterministic string concatenation

Running Example: ToyLLM

为了让后续公式可以被直接计算和实验验证,我们使用一个 finite-horizon episodic MDP ToyLLM 作为示例。 它保留 autoregressive generation 的核心结构,同时将 vocabulary 和最大生成长度控制在可以枚举的范围内。

给定一个固定 prompt xx, 最大 response 长度 TT 和 target response y=(y0,,yL1)y^*=(y_0^*,\ldots,y_{L-1}^*), 其中 LTL\leq T. 对于普通 token 词表 V\overline{\mathcal V}, 我们令 V=V{eos}\mathcal V=\overline{\mathcal V}\cup\{\langle\mathrm{eos}\rangle\}.

时刻 tt 的 state 是 prompt 和已经生成的 prefix, action 是下一个 token:

st=(x,y<t),y<t=(a0,,at1),at=ytV.s_t=(x,y_{<t}),\quad y_{<t}=(a_0,\ldots,a_{t-1}),\quad a_t=y_t\in\mathcal V.

普通 token 会被确定性地追加到 prefix;生成 <eos> 或已经生成 TT 个普通 token 时,episode 终止。 去掉 <eos> 后的完整 response 记为 yy,环境只在 episode 结束时给出 exact-match reward:

R(y)={1,if y=y,0,otherwise.R(y)=\begin{cases} 1,&\text{if }y=y^*,\\ 0,&\text{otherwise}. \end{cases}

对这个固定 prompt, action space 和非终止状态空间分别为

A=V,Snonterminal=t=0T1{x}×Vt.\mathcal A=\mathcal V,\qquad \mathcal S_{\mathrm{non-terminal}}=\bigcup_{t=0}^{T-1}\{x\}\times\overline{\mathcal V}^{\,t}.

K=VK=|\overline{\mathcal V}|. 如果所有结束后的 response 共用一个 terminal state, 则状态总数为

S:=S=1+t=0T1Kt.S:=|\mathcal S|=1+\sum_{t=0}^{T-1}K^t.

因此,即使 transition 是确定性的,prefix state space 仍然随 horizon 指数增长。 真实 LLM 的 vocabulary 通常包含数万到数十万个 token, 使 action space 也非常大。

ToyLLM 还有一个重要性质:初始状态的 value 存在解析解。令成功所需的 action sequence 为

a={(y0,,yL1,eos),L<T,(y0,,yT1),L=T,a^*= \begin{cases} (y_0^*,\ldots,y_{L-1}^*,\langle\mathrm{eos}\rangle),&L<T,\\ (y_0^*,\ldots,y_{T-1}^*),&L=T, \end{cases}

并令 m=am=|a^*|. 如果 reward 在最后一个 action 后产生,则

Vπ(s0)=γm1t=0m1π(atx,a<t).V^\pi(s_0) =\gamma^{m-1}\prod_{t=0}^{m-1} \pi(a_t^*\mid x,a_{<t}^*).

这个解析解可以作为后续 matrix solve, Monte Carlo 和 temporal-difference learning 实现的 ground truth. 若 policy 接近 uniform distribution, 成功概率大致按 Vm|\mathcal V|^{-m} 下降,也直接展示了 terminal verifier reward 带来的 sparse-reward 问题。

ToyLLM 的 exact-match reward 是 verifier 的最简单形式。 真实 RLVR 通常从 prompt distribution xDx\sim\mathcal D 采样问题,并用 Verifier(x,y)\operatorname{Verifier}(x,y) 接受多种等价的正确答案。 ToyLLM 暂时固定一个 prompt 和 target,以便隔离并研究 policy evaluation 与 policy improvement.

Policy Evaluation 一章中,我们将比较依赖完整环境模型的精确解法、只依赖 rollout 的 MC/TD 方法,以及能够在不同 prefix 之间共享信息的 function approximation.

Takeaway

  1. Shen, J., Luo, S., Li, Y., Liu, J., Qu, W., Zhang, Y., Huang, Z., Li, T., Hu, M., Liu, X., Chen, Y., & He, J. (2026). A First-Principles Derivation of LLM Policy Optimization: From Expected Reward to GRPO and Its Structural Extensions. https://arxiv.org/abs/2606.16733

Bellman Equation Theorem

Bellman equation 将 return 的递推结构转化为 value function 的递推结构。 本章我们先研究固定 policy 的 value function, 再研究 optimal policy 中的 optimal value function, 前者通向 policy evaluation, 后者通向 value based methods.

Why Bellman Equation?

在强化学习中,我们关心的是 agent 从当前状态出发,在未来持续决策后能够获得的 expected return. 但是,expected return 依赖于整条未来轨迹:后续会到达什么状态、采取什么动作、获得什么奖励,都是随机的。

直接分析完整 trajectory 的 expected return 往往比较困难。 本节将建立关于 policy value 的基本定理,后续算法可以据此只处理相邻时间步之间的关系。 本节的核心思想在于:

将长期回报写成”当前一步的收益 + 下一状态的未来价值”,把一个全局问题转化为递推问题。 即

long-term return=immediate reward+discounted future return\text{long-term return} = \text{immediate reward} + \text{discounted future return}

Value Function and Q function

首先,我们介绍两个基本概念,分别是 value function Vπ(st)V^{\pi}(s_t) 和 Q function Qπ(st,at)Q^{\pi}(s_t, a_t).

Definition: Value Function

我们定义 value function 如下:

Vπ(st)=Eπ[GtSt=st]V^{\pi}(s_t) = \mathbb{E}^{\pi}[G_t \mid S_t = s_t]

value function 的具体含义为:agent 从当前状态 sts_t 出发,一直遵循当前策略 π\pi,最后获取到的 return.

Definition: Q Function

Q function (state-action value function) 定义如下:

Qπ(st,at)=Eπ[GtSt=st,At=at]Q^{\pi}(s_t, a_t) = \mathbb{E}^{\pi}[G_t \mid S_t=s_t, A_t=a_t]

其具体含义为:agent 从当前状态 sts_t 出发,先执行 action ata_t,再遵循当前策略 π\pi,最后获取到的 return.

为了方便,我们令 Vπ(term)=0V^{\pi}(\langle\text{term}\rangle) = 0, Qπ(term,a)=0,aAQ^{\pi}(\langle\text{term}\rangle, a) = 0, \forall a \in \mathcal{A}.

Value function 和 Q function 存在如下关系:

Proposition: Relationship between value function and Q function
Vπ(st)=Eaπ(st)[Qπ(st,a)]V^{\pi}(s_t) = \mathbb{E}_{a \sim \pi(\cdot \mid s_t)}[Q^{\pi}(s_t, a)]

value function 是 Q function 关于 action ata_t 的期望。

Proof:

由全概率公式,我们有:

Vπ(st)=Eπ[GtSt=st]=aAEπ[GtSt=st,At=a]π(ast)=aAQπ(st,a)π(ast)=Eaπ(st)[Qπ(st,a)]\begin{aligned} V^{\pi}(s_t) &= \mathbb{E}^{\pi}[G_t \mid S_t = s_t] \\ &= \sum_{a \in \mathcal{A}} \mathbb{E}^{\pi}[G_t \mid S_t = s_t, A_t = a] \, \pi(a \mid s_t) \\ &= \sum_{a \in \mathcal{A}} Q^{\pi}(s_t, a) \, \pi(a \mid s_t) = \mathbb{E}_{a \sim \pi(\cdot \mid s_t)}[Q^{\pi}(s_t, a)] \end{aligned}

Value function 和 Q function 存在如下递推关系:

Proposition: Iterative property of Value function and Q function
Vπ(st)=Eatπ(st),(rt,st+1)p(,st,at)[rt+γVπ(st+1)St=st]Qπ(st,at)=E(rt,st+1)p(,st,at),at+1π(st+1)[rt+γQπ(st+1,at+1)St=st,At=at]\boxed{ \begin{aligned} V^{\pi}(s_t) &= \mathbb{E}_{a_t \sim \pi(\cdot \mid s_t), \, (r_t, s_{t+1}) \sim p(\cdot, \cdot \mid s_t, a_t)}[r_t + \gamma V^{\pi}(s_{t+1}) \mid S_t=s_t] \\ Q^{\pi}(s_t, a_t) &= \mathbb{E}_{(r_t, s_{t+1}) \sim p(\cdot, \cdot \mid s_t, a_t), \, a_{t+1} \sim \pi(\cdot \mid s_{t+1})}[r_t + \gamma Q^{\pi}(s_{t+1}, a_{t+1}) \mid S_t=s_t, A_t=a_t] \end{aligned} }
Proof:

我们以 value function 为例:

Vπ(st)=Eπ[GtSt=st]=Eπ[rt+γGt+1St=st]=Eatπ(st),(rt,st+1)p(,st,at)[Eπ[rt+γGt+1st,at,rt,st+1]St=st]=Eatπ(st),(rt,st+1)p(,st,at)[rt+γEπ[Gt+1st+1]St=st]=Eatπ(st),(rt,st+1)p(,st,at)[rt+γVπ(st+1)St=st]\begin{aligned} V^{\pi}(s_t) &= \mathbb{E}^{\pi}[G_t \mid S_t = s_t] \\ &= \mathbb{E}^{\pi}[r_t + \gamma G_{t+1} \mid S_t = s_t] \\ &= \mathbb{E}_{a_t \sim \pi(\cdot \mid s_t), \, (r_t, s_{t+1}) \sim p(\cdot, \cdot \mid s_t, a_t)}\left[\mathbb{E}^{\pi}[r_t + \gamma G_{t+1} \mid s_t, a_t, r_t, s_{t+1}] \mid S_t = s_t\right] \\ &= \mathbb{E}_{a_t \sim \pi(\cdot \mid s_t), \, (r_t, s_{t+1}) \sim p(\cdot, \cdot \mid s_t, a_t)}\left[r_t + \gamma \mathbb{E}^{\pi}[G_{t+1} \mid s_{t+1}] \mid S_t = s_t\right] \\ &= \mathbb{E}_{a_t \sim \pi(\cdot \mid s_t), \, (r_t, s_{t+1}) \sim p(\cdot, \cdot \mid s_t, a_t)}\left[r_t + \gamma V^{\pi}(s_{t+1}) \mid S_t = s_t\right] \end{aligned}

这里第四个等式使用了 Markov property,即 Gt+1G_{t+1} 在已知 st+1s_{t+1} 的条件下与 st,at,rts_t, a_t, r_t 无关。

Q function 的证明类似,略过。

Bellman Equation Theorem

上面的递推关系说明了 value function 和 Q function 满足一个递推关系,接下来, Bellman Equation Theorem 证明了满足这个递归关系的函数一定是 value function/Q function.

Theorem: Bellman Equation Theorem (Value Function)

π\pi 为一个策略,假设 γ(0,1)\gamma \in (0, 1), S<|\mathcal{S}| < \infty 以及 rR<|r| \leq R < \infty, 几乎处处成立 (a.s.a.s.). 那么 π\pi 对应的 value function Vπ:SRV^{\pi}: \mathcal{S} \to \mathbb{R} 存在,且满足 Bellman equation:

Vπ(st)=Eatπ(st),(rt,st+1)p(,st,at)[rt+γVπ(st+1)St=st]\htmlId{BE_value_function}{\begin{equation} V^{\pi}(s_t) = \mathbb{E}_{a_t \sim \pi(\cdot \mid s_t), \, (r_t, s_{t+1}) \sim p(\cdot, \cdot \mid s_t, a_t)}\left[r_t + \gamma V^{\pi}(s_{t+1}) \mid S_t=s_t\right] \end{equation}}

反之,如果存在函数 V:SRV: \mathcal{S} \to \mathbb{R} 满足 Bellman equation, 则 V=VπV = V^{\pi}.

Proof:

iterative property 我们知道 value function 满足 Bellman equation,接下来我们证明唯一性.

假设存在函数 V:SRV: \mathcal{S} \to \mathbb{R} 满足 Bellman equation. 我们定义 Bellman 算子 Tπ\mathcal{T}^{\pi} 为:

(TπV)(st):=Eatπ(st),(rt,st+1)p(,st,at)[rt+γV(st+1)St=st]\htmlId{bellman_operator}{\begin{equation} (\mathcal{T}^{\pi} V)(s_t) := \mathbb{E}_{a_t \sim \pi(\cdot \mid s_t), \, (r_t, s_{t+1}) \sim p(\cdot, \cdot \mid s_t, a_t)}\left[r_t + \gamma V(s_{t+1}) \mid S_t=s_t\right] \end{equation}}

我们证明该算子是一个 contraction mapping, 考虑范数 V=maxsSV(s)\|V\|_{\infty} = \max_{s \in \mathcal{S}} |V(s)|:

(TπV1)(st)(TπV2)(st)=Eatπ(st),(rt,st+1)p(,st,at)[γV1(st+1)γV2(st+1)]γEatπ(st),(r,st+1)p(,st,at)V1(st+1)V2(st+1)γmaxsSV1(s)V2(s)=γV1V2\begin{aligned} \left|(\mathcal{T}^{\pi} V_1)(s_t) - (\mathcal{T}^{\pi} V_2)(s_t)\right| &= \left|\mathbb{E}_{a_t \sim \pi(\cdot \mid s_t), \, (r_t, s_{t+1}) \sim p(\cdot, \cdot \mid s_t, a_t)}\left[\gamma V_1(s_{t+1}) - \gamma V_2(s_{t+1})\right]\right| \\ &\leq \gamma \, \mathbb{E}_{a_t \sim \pi(\cdot \mid s_t), \, (r, s_{t+1}) \sim p(\cdot, \cdot \mid s_t, a_t)}\left|V_1(s_{t+1}) - V_2(s_{t+1})\right| \\ &\leq \gamma \max_{s \in \mathcal{S}}\left|V_1(s) - V_2(s)\right| \\ &= \gamma \left\|V_1 - V_2\right\|_{\infty} \end{aligned}

上式对任意 sSs \in \mathcal{S} 成立,因此:

TπV1TπV2γV1V2\left\|\mathcal{T}^{\pi} V_1 - \mathcal{T}^{\pi} V_2\right\|_{\infty} \leq \gamma \left\|V_1 - V_2\right\|_{\infty}

由于 γ<1\gamma < 1, 因此 Tπ\mathcal{T}^{\pi} 是一个 contraction mapping.

根据 不动点定理,已知 VπV^{\pi} 是一个不动点(满足 Bellman equation),而不动点唯一,因此 V=VπV = V^{\pi}.

对于 Q function, 我们也有相同的结论:

Theorem: Bellman Equation Theorem (Q Function)

π\pi 为一个策略,假设 γ(0,1)\gamma \in (0, 1), S<|\mathcal{S}| < \infty, A<|\mathcal{A}| < \infty 以及 rR<|r| \leq R < \infty, 几乎处处成立 (a.s.a.s.). 那么 π\pi 对应的 Q function Qπ:S×ARQ^{\pi}: \mathcal{S} \times \mathcal{A} \to \mathbb{R} 存在,且满足 Bellman equation:

Qπ(st,at)=E(rt,st+1)p(,st,at),at+1π(st+1)[rt+γQπ(st+1,at+1)St=st,At=at]\boxed{ Q^{\pi}(s_t, a_t) = \mathbb{E}_{(r_t, s_{t+1}) \sim p(\cdot, \cdot \mid s_t, a_t), \, a_{t+1} \sim \pi(\cdot \mid s_{t+1})}\left[r_t + \gamma Q^{\pi}(s_{t+1}, a_{t+1}) \mid S_t=s_t, A_t=a_t\right] }

反之,如果存在函数 Q:S×ARQ: \mathcal{S} \times \mathcal{A} \to \mathbb{R} 满足 Bellman equation, 则 Q=QπQ = Q^{\pi}.

Proof:

证明与 value function 的证明类似,略过。

Bellman Optimality Equation Theorem

上一节我们介绍了针对一般 policy 的 Bellman equation, 这一节我们给出针对 optimal policy 的 Bellman Optimal Equation.

我们首先给出 optimal policy 的定义

Definition: Optimal Policy

如果策略 π\pi^* 满足:

Vπ(s)Vπ(s),sS,πV^{\pi^*}(s) \geq V^{\pi}(s), \quad \forall s \in \mathcal{S}, \, \forall \pi

则称策略 π\pi^*optimal policy. 对应的 VπV^{\pi^*}QπQ^{\pi^*} 分别称为 optimal value functionoptimal Q function, 简记为 V=VπV^* = V^{\pi^*}, Q=QπQ^* = Q^{\pi^*}.

optimal policy 通常依赖于状态 ss,并且可能不唯一;但所有 optimal policy 对应同一个 optimal value function VV^*.

与 value function VπV^\pi 和 Q function QπQ^\pi 之间存在联系一样,optimal value function VV^{*} 和 optimal Q function QQ^{*} 也存在类似关系

Proposition: Relationship between optimal value function and optimal Q function
V(st)=maxaAQ(st,a)\boxed{V^*(s_t) = \max_{a \in \mathcal{A}} Q^*(s_t, a)}

VV^* 就是 QQ^* 的最优动作对应的 return.

Proof:

我们先证左边小于等于右边:

V(st)=Eaπ(st)[Q(st,a)]Eaπ(st)[maxaAQ(st,a)]=maxaAQ(st,a)V^*(s_t) = \mathbb{E}_{a \sim \pi^*(\cdot \mid s_t)}[Q^*(s_t, a)] \leq \mathbb{E}_{a \sim \pi^*(\cdot \mid s_t)}[\max_{a \in \mathcal{A}} Q^*(s_t, a)] = \max_{a \in \mathcal{A}} Q^*(s_t, a)

再证右边小于等于左边:

aargmaxaQ(st,a)a^* \in \arg\max_{a} Q^*(s_t, a), 令策略 π\pi' 为确定性策略 π(ast)=1\pi'(a^* \mid s_t) = 1, 则:

maxaAQ(st,a)=Qπ(st,a)=Vπ(st)V(st)\max_{a \in \mathcal{A}} Q^*(s_t, a^*) = Q^{\pi'}(s_t, a^*) = V^{\pi'}(s_t) \leq V^*(s_t)

因此两边相等。

接下来,我们给出关于 optimal value function VV^* 对应的 Bellman optimality equation.

Theorem: Bellman Optimality Equation Theorem (value function)

假设 γ(0,1)\gamma \in (0, 1), S<|\mathcal{S}| < \infty, A<|\mathcal{A}| < \infty 以及 rR<|r| \leq R < \infty, 几乎处处成立 (a.s.). 那么 optimal value function V:SRV^*:\mathcal{S} \to \mathbb{R} 存在,且满足 Bellman optimality equation:

V(st)=maxaAE(rt,st+1)p(,st,at)[rt+γV(st+1)St=st,At=at]\boxed{ V^*(s_t) = \max_{a \in \mathcal{A}} \mathbb{E}_{(r_t, s_{t+1}) \sim p(\cdot, \cdot \mid s_t, a_t)}\left[r_t + \gamma V^*(s_{t+1}) \mid S_t=s_t,A_t=a_t\right] }

反之,如果存在函数 V:SRV: \mathcal{S} \to \mathbb{R} 满足 Bellman optimality equation, 则 V=VV = V^*.

最后,我们可以给出一个 optimal deterministic policy:

π(st)=argmaxaAE(rt,st+1)p(,st,at)[rt+γV(st+1)St=st,At=at]=argmaxaAQ(st,a)\htmlId{optimal_deterministic_policy}{\begin{align} \pi^*(s_t) &= \arg\max_{a \in \mathcal{A}} \mathbb{E}_{(r_t, s_{t+1}) \sim p(\cdot, \cdot \mid s_t, a_t)}\left[r_t + \gamma V^*(s_{t+1}) \mid S_t=s_t, A_t=a_t\right]\\ &= \arg\max_{a \in \mathcal{A}} Q^*(s_t, a) \end{align}}

证明需要用到下面的引理

Lemma: Monotonicity of Bellman operator

π\pi 为一个策略,Tπ\mathcal{T}^{\pi}T\mathcal{T}^* 分别是 Bellman operator 和 Bellman optimality operator, 其中 Bellman optimality operator 定义为

(TV)(st):=maxaAE(rt,st+1)p(,st,at)[rt+γV(st+1)St=st,At=at]\htmlId{bellman_optimality_operator}{\begin{equation} (\mathcal{T}^* V)(s_t) := \max_{a \in \mathcal{A}} \mathbb{E}_{(r_t, s_{t+1}) \sim p(\cdot, \cdot \mid s_t, a_t)}\left[r_t + \gamma V(s_{t+1}) \mid S_t=s_t, A_t=a_t\right] \end{equation}}

我们有:

  1. 对任意 V:SRV: \mathcal{S} \to \mathbb{R}, 有 TπV(s)TV(s),sS\mathcal{T}^{\pi} V(s) \leq \mathcal{T}^* V(s), \, \forall s \in \mathcal{S}.
  2. 对任意 U,V:SRU, V: \mathcal{S} \to \mathbb{R}, 如果 UVU \leq V, 则 TU(s)TV(s),sS\mathcal{T}^* U(s) \leq \mathcal{T}^* V(s), \, \forall s \in \mathcal{S}.

证明如下

Proof:

我们首先证明 Bellman optimality operator T\mathcal{T}^* 是一个 contraction mapping: 对任意 V1,V2V_1, V_2:

(TV1)(st)(TV2)(st)=maxaE[rt+γV1(st+1)]maxaE[rt+γV2(st+1)]maxaE[rt+γV1(st+1)]E[rt+γV2(st+1)]=γmaxaE[V1(st+1)V2(st+1)]γV1V2\begin{aligned} \left|(\mathcal{T}^* V_1)(s_t) - (\mathcal{T}^* V_2)(s_t)\right| &= \left|\max_{a} \mathbb{E}[r_t + \gamma V_1(s_{t+1})] - \max_{a} \mathbb{E}[r_t + \gamma V_2(s_{t+1})]\right| \\ &\leq \max_{a} \left|\mathbb{E}[r_t + \gamma V_1(s_{t+1})] - \mathbb{E}[r_t + \gamma V_2(s_{t+1})]\right| \\ &= \gamma \max_{a} \left|\mathbb{E}[V_1(s_{t+1}) - V_2(s_{t+1})]\right| \\ &\leq \gamma \|V_1 - V_2\|_{\infty} \end{aligned}

这里第一个不等式使用了 maxsv(s)maxsu(s)maxsu(s)v(s)|\max_s v(s) - \max_s u(s)| \leq \max_s |u(s) - v(s)|.

因此 TV1TV2γV1V2\|\mathcal{T}^* V_1 - \mathcal{T}^* V_2\|_{\infty} \leq \gamma \|V_1 - V_2\|_{\infty}, T\mathcal{T}^* 是 contraction mapping.

根据不动点定理T\mathcal{T}^* 存在唯一不动点 VV^*.

我们定义:

π(st)=argmaxaAE(rt,st+1)p(,st,at)[r+γV(st+1)St=st,At=at]\pi^*(s_t) = \arg\max_{a \in \mathcal{A}} \mathbb{E}_{(r_t, s_{t+1}) \sim p(\cdot, \cdot \mid s_t, a_t)}\left[r + \gamma V^*(s_{t+1}) \mid S_t=s_t, A_t=a_t\right]

此时,我们有 TODO

V(st)=Eπ[rt+γV(st+1)st]=Eπ[rt+γEπ[rt+1+γV(st+2)st+1]st]=Eπ[Eπ[rt+γ(rt+1+γV(st+2))st+1]st]=Eπ[rt+γrt+1+γ2V(st+2)st]=Eπ[rt+γrt+1+γ2rt+2+st]=Vπ(st)\begin{aligned} V^*(s_t) &= \mathbb{E}^{\pi^*}\left[r_t+\gamma V^*(s_{t+1})\mid s_t\right]\\ &= \mathbb{E}^{\pi^*}\left[r_t+\gamma \mathbb{E}^{\pi^*}\left[r_{t+1}+\gamma V^*(s_{t+2})\mid s_{t+1}\right]\mid s_t\right]\\ &= \mathbb{E}^{\pi^*}\left[\mathbb{E}^{\pi^*}\left[r_t+\gamma (r_{t+1}+\gamma V^*(s_{t+2}))\mid s_{t+1}\right]\mid s_t\right]\\ &= \mathbb{E}^{\pi^*}\left[r_t+\gamma r_{t+1}+\gamma^2 V^*(s_{t+2})\mid s_t\right]\\ &= \mathbb{E}^{\pi^*}\left[r_t+\gamma r_{t+1}+\gamma^2r_{t+2}+\cdots\mid s_t\right]= V^{\pi^*}(s_t) \end{aligned}

V=VπV^*=V^{\pi^*}, 这里第四个等式使用了 Law of Total Expectation.

现在我们证明 π\pi^* 是最优策略,令 π\pi 为任意一个策略,我们有

Vπ=TπVTV(T)2(Vπ)(T)k(Vπ)kfixed-point / contractionV.V^\pi=\mathcal{T}^{\pi}V \le \mathcal{T}^*V \leq (\mathcal{T}^*)^2(V^\pi) \leq \cdots \leq (\mathcal{T}^*)^k(V^\pi) \xrightarrow[k\to\infty]{\text{fixed-point / contraction}} V^*.

VπVV^\pi\leq V^*, 这里第一个等式和第二个等式利用了 Lemma, 由于 VπV^\pi 是任意的,我们有 V=VπV^*=V^{\pi^*}.

对于 optimal Q function QQ^*, 我们也有类似的结论。

Theorem: Bellman Optimality Equation Theorem (Q function)

假设 γ(0,1)\gamma \in (0, 1), S<|\mathcal{S}| < \infty, A<|\mathcal{A}| < \infty 以及 rR<|r| \leq R < \infty, 几乎处处成立 (a.s.). 那么 optimal Q function Q:S×ARQ^*: \mathcal{S} \times \mathcal{A} \to \mathbb{R} 存在,且满足 Bellman optimality equation:

Q(st,at)=E(rt,st+1)p(,st,at)[r+γmaxat+1AQ(st+1,at+1)St=st,At=at]\boxed{ Q^*(s_t, a_t) = \mathbb{E}_{(r_t, s_{t+1}) \sim p(\cdot, \cdot \mid s_t, a_t)}\left[r + \gamma \max_{a_{t+1} \in \mathcal{A}} Q^*(s_{t+1}, a_{t+1}) \mid S_t=s_t, A_t=a_t\right] }

反之,如果存在函数 Q:S×ARQ: \mathcal{S} \times \mathcal{A} \to \mathbb{R} 满足 Bellman optimality equation, 则 Q=QQ = Q^*.

Proof:

证明与 value function 类似,略过。

下一章我们将从这个精确但不可扩展的起点出发,逐步引入 sampling 和 function approximation 来完成 policy evaluation.

Takeaway

  1. Gt=rt+γGt+1G_t=r_t+\gamma G_{t+1} 使 return 可以递归分解;对它取条件期望便得到 Bellman equation.
  2. Bellman equation 回答“固定 policy 的价值是什么”,Bellman optimality equation 回答“所有 policy 中的最优价值是什么”。
  3. 对 LLM,value 描述从 (prompt, token prefix) 继续生成后的期望最终得分; Bellman recursion 则承担将稀疏 terminal reward 传回早期 token 的 credit assignment.

Policy Evaluation

Bellman Equation Theorem 一节中,我们看到 VπV^\pi 是 Bellman operator Tπ\mathcal T^\pi 的唯一不动点。 Bellman optimality equation 直接寻找最优价值,而本章先解决一个更基础的 evaluation 问题:

给定一个策略 π\pi,如何高效求解策略 π\pi 对应的 value function VπV^\pi?

我们将先从完整 MDP 模型的精确解法,逐步过渡到只依赖 rollout 的 MC/TD 方法和现代 RL 使用的函数逼近方法。 这条路径同时回答两个问题:

  1. 为什么 Bellman equation 在小环境中可直接求解
  2. 精确求解方法对于 LLM 来说存在什么问题

Tabular Learning

这一节我们介绍针对离散小空间场景下的求解和优化方案。

Tabular Learning with Matrix Solve

对于离散状态空间来说,我们可以将 value iteration 算法中的 Bellman Equation 写为如下的矩阵形式

Vπ=Rπ+γPπVπV_{\pi} = R_{\pi} + \gamma P_{\pi} V_{\pi}

其中

Vπ=[Vπ(s1),,Vπ(sS)]TRSRπ=[Rπ(s1),,Rπ(sS)]TRSPπ=[Pπ(sjsi)]RS×S\begin{aligned} V_\pi &= [V_{\pi}(s_1),\dots, V_{\pi}(s_S)]^T\in\mathbb{R}^S\\ R_\pi &=[R_{\pi}(s_1),\dots, R_{\pi}(s_S)]^T\in\mathbb{R}^S\\ P_\pi &=[P_\pi(s_j\mid s_i)]\in\mathbb{R}^{S\times S} \end{aligned}

可以看到,上面这个其实是一个线性方程组,即

(ISγPπ)Vπ=Rπ(\mathbf{I}_S - \gamma P_{\pi})V_{\pi} = R_{\pi}

对于这个线性方程组,我们可以用矩阵求逆或者迭代解法求解。 精确算法如下所示

Algorithm: Value function Matrix Exact Solve

Input: policy π\pi, discount factor γ\gamma, state space S\mathcal{S} with size SS

  1. assign index i(s)i(s) for sS,i=1,,Ss\in\mathcal{S}, i=1,\dots,S
  2. P=0S×SP=\mathbf{0}_{S\times S}, R=0SR=\mathbf{0}_{S}
  3. For each sSs\in\mathcal{S}
    1. ii(s)i\gets i(s)
    2. For each aAa\in\mathcal{A}
      • (s,r,done)STEP(s,a)(s', r, \mathrm{done}) \gets \mathrm{STEP}(s, a)
      • R[i]R[i]+π(as)rR[i]\gets R[i]+\pi(a\mid s)\cdot r
      • If NOT done\mathrm{done}
        • ji(s)j\gets i(s')
        • P[i,j]P[i,j]+π(as)P[i,j]\gets P[i,j]+\pi(a\mid s)
  4. VSOLVE(ISγP,R)V\gets \mathrm{SOLVE}(I_S-\gamma P, R)
  5. return VV

迭代算法则是给定随机初始化的 V0V_0, 我们通过迭代计算得到序列 {Vk}\{V_k\}, 其中 Vk+1=Rπ+γPπVkV_{k+1}=R_{\pi} + \gamma P_{\pi}V_k, 我们可以证明这个序列收敛到 VπV_{\pi}.

Algorithm: Value function Matrix Iterative Solve

Input: policy π\pi, discount factor γ\gamma, convergence threshold θ\theta, state space S\mathcal{S}

  1. V(s)=0,sSV(s)=0, s\in\mathcal{S}
  2. While Δ>θ\Delta > \theta:
    1. Δ0\Delta\gets0
    2. For each sSs\in\mathcal{S}
      1. v0v\gets0
      2. For each aAa\in\mathcal{A}
        • (s,r,done)STEP(s,a)(s', r, \mathrm{done}) \gets \mathrm{STEP}(s, a)
        • If done\mathrm{done}
          • vv+π(as)rv \gets v + \pi(a|s) \cdot r
        • Else
          • vv+π(as)(r+γV(s))v \gets v + \pi(a|s) \cdot (r + \gamma \cdot V(s'))
      3. Δmax(Δ,vV(s))\Delta ← max(\Delta, |v − V(s)|)
      4. V(s)vV(s) \gets v
  3. return VV

下面是两种算法在不同规模下对应的求解时间变化曲线

Comparison of matrix solve and iterative solve

可以看到,由于稠密线性方程组的直接求解通常需要 O(S3)\mathcal{O}(S^3) 时间和 O(S2)\mathcal{O}(S^2) 空间,随着状态数 SS 的增加,其所需时间急剧上升。 而迭代解法虽然比矩阵直接求解效率更高,但是仍然需要枚举全部状态,当状态空间增大之后,迭代解法的时间也无法接受。 因此,我们需要计算更加高效的算法。

Tabular Learning with MC

VπV^{\pi} 定义,我们有:

Vπ(s0)=Eπ[R(τ)s0]V^{\pi}(s_0) = \mathbb{E}^{\pi}\left[R(\tau) \mid s_0\right]

由于 return 是一个随机变量,我们可以利用 Monte Carlo (MC) 方法得到一个 Unbiased Estimator,即从 s0s_0 出发,独立随机采样 (i.i.d.) MM 条轨迹:

{τ(i)=(s0,a0(i),r0(i),,aT(i)1(i),rT(i)1(i),sT(i)(i))},i=1,,M\{\tau^{(i)} = (s_0, a_0^{(i)}, r_0^{(i)}, \dots, a_{T^{(i)}-1}^{(i)}, r_{T^{(i)}-1}^{(i)}, s_{T^{(i)}}^{(i)})\}, \quad i = 1, \dots, M

然后我们使用样本平均 (sample mean) 来近似期望:

Vπ(s0)=Eπ[R(τ)s0]1Mi=1MR(τ(i))V^{\pi}(s_0) = \mathbb{E}^{\pi}[R(\tau) \mid s_0] \approx \frac{1}{M} \sum_{i=1}^{M} R(\tau^{(i)})

对应的算法如下所示

Algorithm: Tabular Learning with MC

Input: policy π\pi, discount factor γ\gamma, episodes MM, state space S\mathcal{S}, initial state distribution p0p_0

  1. V(s)=0,ReturnSum(s)=0,ReturnCount(s)=0,sSV(s)=0, \mathrm{ReturnSum}(s)=0, \mathrm{ReturnCount}(s)=0, s\in\mathcal{S}
  2. For episode=1,2,,M\mathrm{episode}=1,2,\dots,M:
    1. s0p0s_0\sim p_0, τ=[]\tau=[]
    2. While not done\mathrm{done}:
      • aπ(s)a\sim\pi(\cdot\mid s)
      • (s,r,done)STEP(s,a)(s', r, \mathrm{done}) \gets \mathrm{STEP}(s, a)
      • τ.append((s,r))\tau.\mathrm{append}((s,r))
      • sss\gets s'
    3. Returns=[]\mathrm{Returns}=[], G0G\gets0
    4. For (s,r)(s,r) in reversed(τ)\mathrm{reversed}(\tau)
      • Gr+γGG\gets r + \gamma G
      • Returns.prepend((s,G))\mathrm{Returns}.\mathrm{prepend}((s,G))
    5. Visited=\mathrm{Visited}=\emptyset
    6. For (s,G)(s,G) in Returns\mathrm{Returns}
      • If s∉Visiteds\not\in \mathrm{Visited}
        • Visited.add(s)\mathrm{Visited}.\mathrm{add}(s)
        • ReturnSum(s)ReturnSum(s)+G\mathrm{ReturnSum}(s) \gets \mathrm{ReturnSum}(s) +G
        • ReturnCount(s)ReturnCount(s)+1\mathrm{ReturnCount}(s) \gets \mathrm{ReturnCount}(s)+ 1
        • V(s)ReturnSum(s)/ReturnCount(s)V(s)\gets \mathrm{ReturnSum}(s) / \mathrm{ReturnCount}(s)
  3. return VV

上面的伪代码使用了 First-Visit MC:在一个 episode 内,同一个状态可能出现多次,但我们只用它第一次出现位置之后的 return 来更新该状态。

Tabular Learning with TD

MC 的优点是 estimator 是一个 unbiased estimator, 其缺点在于我们必须完整采样完一条轨迹才能得到这个估计,对于 long-horizon task, MC 方法效率很低。

为了解决这个问题,我们可以使用 Temporal difference (TD) learning 方法。 Temporal difference (TD) learning 通过相邻两步的转换关系来近似 value function. 其核心思想在于

通过新的信息和新的预测来更新旧的预测

注意到:

Vπ(st)=Eπ[rt+γVπ(st+1)St=st]V^{\pi}(s_t) = \mathbb{E}^{\pi}[r_t + \gamma V^{\pi}(s_{t+1}) \mid S_t=s_t]

我们首先定义 TD target 和 TD error 如下

Definition: TD target and TD error

给定当前 estimator VtV_t,我们定义 one-step TD target

Vt={rt,if st+1 is terminal,rt+γVt(st+1),otherwise.\overline{V}_t = \begin{cases} r_t,&\text{if }s_{t+1}\text{ is terminal},\\ r_t+\gamma V_t(s_{t+1}),&\text{otherwise}. \end{cases}

定义 TD error

δt=VtVt(st).\delta_t = \overline{V}_t - V_t(s_t).

接下来,我们基于 TD error 更新 VtV_t

Vt+1(s)={Vt(st)+αt(st)δt,if s=st,Vt(s),otherwise.V_{t+1}(s) = \begin{cases} V_t(s_t) + \alpha_t(s_t)\delta_t,&\text{if }s=s_t,\\ V_t(s),&\text{otherwise}. \end{cases}

其中 αt(st)(0,1)\alpha_t(s_t) \in (0, 1) 是 learning rate.

TD learning 对应的算法如下所示

Algorithm: Tabular Learning with TD

Input: policy π\pi, discount factor γ\gamma, learning rate α\alpha, episodes MM, state space S\mathcal{S}, initial state distribution p0p_0

  1. V(s)=0,sSV(s)=0, s\in\mathcal{S}
  2. For episode=1,2,,M\mathrm{episode}=1,2,\dots,M:
    1. s0p0s_0\sim p_0
    2. While not done\mathrm{done}:
      • aπ(s)a\sim\pi(\cdot\mid s)
      • (s,r,done)STEP(s,a)(s', r, \mathrm{done}) \gets \mathrm{STEP}(s, a)
      • If done\mathrm{done}
        • yry\gets r
      • Else
        • yr+γV(s)y \gets r + \gamma \cdot V(s')
      • V(s)V(s)+α(yV(s))V(s)\gets V(s) + \alpha\cdot(y-V(s))
      • sss\gets s'
  3. return VV

接下来我们来分析一下算法的正确性。 首先,Vt\overline{V}_t 被称为 TD target 的原因为

Vt+1(st)Vt(st)=Vt(st)αt(st)δtVt(st)=1αt(st)Vt(st)Vt(st)<Vt(st)Vt(st)\begin{aligned} |V_{t+1}(s_t)-\overline{V}_t(s_t)| &= |V_t(s_t) - \alpha_t(s_t)\delta_t-\overline{V}_t(s_t)|\\ &= |1-\alpha_t(s_t)||V_{t}(s_t)-\overline{V}_t(s_t)|\\ &< |V_{t}(s_t)-\overline{V}_t(s_t)| \end{aligned}

可以看到函数序列 {Vt}\{V_t\} 收敛到 Vt\overline{V}_t. 另一方面,我们有

E[δtSt=st]=E[Vπ(St)(rt+γVπ(st+1))St=st]=Vπ(st)E[rt+γVπ(st+1)St=st]=0\begin{aligned} \mathbb{E}[\delta_t\mid S_t=s_t] &= \mathbb{E}[V^{\pi}(S_t)-\left(r_t + \gamma V^{\pi}(s_{t+1})\right)\mid S_t=s_t]\\ &= V^{\pi}(s_t) - \mathbb{E}[r_t + \gamma V^{\pi}(s_{t+1})\mid S_t=s_t]\\ &=0 \end{aligned}

因此 TD error δt\delta_t 衡量了当前的估计 VtV_tVπV^\pi 之间的差距。

算法的正确性由如下定理给定

Theorem: Convergence of TD Learning

给定策略 π\pi, 基于 TD learning algorithm, 对任意 sSs\in\mathcal{S}, 如果 tαt(s)=\sum_t\alpha_t(s)=\infty, tαt2(s)<\sum_t\alpha_t^2(s)<\infty, 则 Vt(s)V_t(s) a.s.a.s. 收敛到 Vπ(s)V^{\pi}(s), tt\to\infty.

Proof:

See Theorem 7.1 of Mathematical Foundations of Reinforcement Learning for more details.

Tabular Learning Experiments

我们在不同 vocab size 下对比 matrix solve, MC 和 TD. 对于 MC 和 TD, 我们将 episodes 设置为 10,00010,000, TD 的 learning rate 设置为 0.10.1. 结果如下图所示

Comparison of Exact solution and MC/TD

可以看到,matrix exact/iterative solve 每次都需要枚举状态,而 MC/TD 的单条数据成本主要由 rollout 长度决定,因此我们可以在不枚举整个状态空间的情况下更新被访问的状态。 不过,固定 rollout budget 并不保证所有状态都得到充分访问,因此随着状态空间的增大,MC/TD 方法的误差也逐渐增高。

我们采样成功的概率为

psuccess=t=0m1π(ata<t).p_{\mathrm{success}}=\prod_{t=0}^{m-1}\pi(a_t^*\mid a_{<t}^*).

如果 policy 接近均匀分布,该概率大致按 Vm|\mathcal V|^{-m} 下降。 因此,我们额外统计不同 vocab size 下,MC 和 TD 达到指定误差所需的 episodes. 我们将 maximum episodes 设置为 100,000100,000,目标 error 设置为 0.010.01,结果如下:

Number of rollouts required to reach the target error

可以看到,在当前 terminal binary reward、固定超参数和误差指标下,MC 比 TD 更快达到目标误差。 当 vocab size 超过 30 后,两者在 rollout budget 内都很难观察到足够多的正样本。

这里暴露出两个不同的问题:

  1. tabular learning 方法不能在相似 prefix 之间共享信息
  2. 稀疏 verifier reward 导致有效样本极少。

函数逼近主要缓解前一个问题;后一个问题还需要更好的 exploration, reward shaping 或更强的初始 policy.

Function Approximation

上面我们介绍了 MC 和 TD 两种 value function 估计方法,但是当动作空间和状态空间特别大时,tabular learning 仍然存在扩展性问题。 为了解决这个扩展性不足的问题,当状态空间 S\mathcal{S} 较大或连续时,我们可以使用函数逼近的方法来近似 VπV^{\pi}.

这里我们使用 Vϕ:SRV_{\phi}:\mathcal{S}\to\mathbb{R} 近似 VπV^{\pi},其中 ϕ\phi 是参数。 令 dπd^\pi 表示策略 π\pi 在 rollout 中诱导的状态访问分布,我们使用 mean squared value (MSE) error:

L(ϕ)=Esdπ[12(Vϕ(s)Vπ(s))2].\mathcal{L}(\phi) = \mathbb{E}_{s \sim d^\pi}\left[\frac{1}{2}\left(V_{\phi}(s) - V^{\pi}(s)\right)^2\right].

Vπ(s)V^\pi(s) 满足 Vπ(s)=Eπ[GS=s]V^\pi(s)=\mathbb E^\pi[G\mid S=s], 因此:

ϕL(ϕ)=Esdπ[(Vϕ(s)Vπ(s))ϕVϕ(s)]=Esdπ,Gπ(s)[(Vϕ(s)G)ϕVϕ(s)].\begin{aligned} \nabla_{\phi} \mathcal{L}(\phi) &= \mathbb{E}_{s \sim d^\pi} \left[\left(V_{\phi}(s)-V^\pi(s)\right)\nabla_{\phi}V_{\phi}(s)\right]\\ &= \mathbb{E}_{s \sim d^\pi,\,G\sim\pi(\cdot\mid s)} \left[\left(V_{\phi}(s)-G\right)\nabla_{\phi}V_{\phi}(s)\right]. \end{aligned}

Function Approximation with MC

现在我们可以使用 MC 方法来估计梯度:

ϕL(ϕ)g:=1Mi=1M(Vϕ(s(i))G(i))ϕVϕ(s(i)).\nabla_{\phi} \mathcal{L}(\phi) \approx g := \frac{1}{M}\sum_{i=1}^{M} \left(V_{\phi}(s^{(i)})-G^{(i)}\right)\nabla_{\phi}V_{\phi}(s^{(i)}).

结合 MC 和 SGD 的算法如下:

Algorithm: Value function approximation with MC

while not converged:

  1. s0p0s_0 \sim p_0, t=0t = 0, τ=[]\tau=[].
  2. while stterms_t \neq \langle\text{term}\rangle:
    • atπ(st)a_t \sim \pi(\cdot \mid s_t)
    • (rt,st+1)p(,st,at)(r_t, s_{t+1}) \sim p(\cdot, \cdot \mid s_t, a_t)
    • τ.append((st,rt))\tau.\mathrm{append}((s_t,r_t))
    • tt+1t \gets t + 1
  3. Set T=tT = t.
  4. For every t=0,,T1t=0,\ldots,T-1, compute Gt=k=tT1γktrkG_t=\sum_{k=t}^{T-1}\gamma^{k-t}r_k.
  5. g1Tt=0T1(Vϕ(st)Gt)ϕVϕ(st)g\gets\frac{1}{T}\sum_{t=0}^{T-1} (V_\phi(s_t)-G_t)\nabla_\phi V_\phi(s_t).
  6. Update ϕ\phi using gg with an optimizer.

Function Approximation with TD

如果能够访问真实的 Vπ(s)V^\pi(s'),我们可以将一步使用 TD target 当作监督信号:

L(ϕ)=Esdππ[12(Vϕ(s)(r+γVπ(s)))2].\mathcal{L}(\phi) = \mathbb{E}_{s \sim d^\pi}^\pi \left[\frac{1}{2}\left(V_{\phi}(s)-\left(\textcolor{red}{r+\gamma V^{\pi}(s')}\right)\right)^2\right].

对应的梯度为:

ϕL(ϕ)=Esdππ[(Vϕ(s)rγVπ(s))ϕVϕ(s)].\nabla_{\phi} \mathcal{L}(\phi) =\mathbb{E}_{s \sim d^\pi}^{\pi} \left[\left(V_{\phi}(s)-r-\gamma \textcolor{red}{V^{\pi}(s')}\right) \nabla_{\phi}V_{\phi}(s)\right].

这里梯度包含了一个未知量 Vπ(s)V^{\pi}(s'). 为了解决这个问题,一个自然的想法是使用当前的 value function Vϕ(s)V_{\phi}(s') 来进行代替:

g:=Esdππ[(Vϕ(s)rγVϕ(s))ϕVϕ(s)].g:=\mathbb{E}_{s \sim d^\pi}^{\pi} \left[\left(V_{\phi}(s)-r-\gamma V_{\phi}(s')\right) \nabla_{\phi}V_{\phi}(s)\right].

对应的数学更新为:

Algorithm: Value function approximation with TD (not correct)

While not converged:

  1. s0p0s_0 \sim p_0, a0π(s0)a_0 \sim \pi(\cdot \mid s_0), (r0,s1)p(,s0,a0)(r_0, s_1) \sim p(\cdot, \cdot \mid s_0, a_0).
  2. compute TD target y={r0+γVϕ(s1),if s1termr0,otherwisey = \begin{cases} r_0 + \gamma V_{\phi}(s_1), & \text{if } s_1 \neq \langle\text{term}\rangle \\ r_0, & \text{otherwise} \end{cases}
  3. g=(Vϕ(s0)y)ϕVϕ(s0)g = (V_{\phi}(s_0) - y) \nabla_{\phi} V_{\phi}(s_0).
  4. Update ϕ\phi using gg with an optimizer.

我们来看一下上述算法对应的 PyTorch 实现:

pred = V_phi(s0)
target = r + gamma * V_phi(s1)
td_error = pred - target
loss = 0.5 * td_error ** 2
loss.backward()

可以看到,我们实际上求的梯度是:

ϕ[12(Vϕ(s)(r+γVϕ(s)))2]=(Vϕ(s)(r+γVϕ(s)))(ϕVϕ(s)γϕVϕ(s))g\begin{aligned} &\nabla_{\phi}\left[\frac{1}{2}\left(V_{\phi}(s) - (r + \gamma V_{\phi}(s'))\right)^2\right] \\ &= (V_{\phi}(s) - (r + \gamma V_{\phi}(s'))) (\nabla_{\phi} V_{\phi}(s) - \gamma\nabla_{\phi} V_{\phi}(s'))\\ &\neq g \end{aligned}

为了解决这个问题,我们可以使用 stop-gradient 技巧来避免 Vϕ(s)V_{\phi}(s') 参与反向传播,此时目标函数变为:

L(ϕ)=12(Vϕ(s)(r+γsg[Vϕ(s)]))2\mathcal{L}(\phi)=\frac{1}{2}\left(V_{\phi}(s) - (r + \gamma \, \mathrm{sg}[V_{\phi}(s')])\right)^2

其中 sg[]\mathrm{sg}[\cdot] 是 stop-gradient operator, 满足:

sg[x]={xforward pass0backward pass\mathrm{sg}[x] = \begin{cases} x & \text{forward pass} \\ 0 & \text{backward pass} \end{cases}

对应的梯度就是:

ϕ[12(Vϕ(s)(r+γsg[Vϕ(s)]))2]=(Vϕ(s)(r+γVϕ(s)))ϕVϕ(s)=g\nabla_{\phi}\left[\frac{1}{2}\left(V_{\phi}(s) - (r + \gamma \, \mathrm{sg}[V_{\phi}(s')])\right)^2\right] = (V_{\phi}(s) - (r + \gamma V_{\phi}(s')))\nabla_{\phi} V_{\phi}(s) = g

对应的 Python 代码为:

pred = V_phi(s0)
target = r + gamma * V_phi(s1)
td_error = pred - target.detach()  # stop-gradient operator
loss = 0.5 * td_error ** 2
loss.backward()

最终,正确的 TD 算法如下:

Algorithm: Value function approximation with TD

While not converged:

  1. s0p0s_0 \sim p_0, a0π(s0)a_0 \sim \pi(\cdot \mid s_0), (r0,s1)p(,s0,a0)(r_0, s_1) \sim p(\cdot, \cdot \mid s_0, a_0).
  2. compute TD target y={r0+γsg[Vϕ(s1)],if s1termr0,otherwisey = \begin{cases} r_0 + \gamma \, \mathrm{sg}[V_{\phi}(s_1)], & \text{if } s_1 \neq \langle\text{term}\rangle \\ r_0, & \text{otherwise} \end{cases}
  3. g=(Vϕ(s0)y)ϕVϕ(s0)g = (V_{\phi}(s_0) - y) \nabla_{\phi} V_{\phi}(s_0).
  4. Update ϕ\phi using gg with an optimizer.

上述做法是 semi-gradient method:bootstrap target 依赖当前参数,但更新时不对 target 求导。因此它一般不是某个固定监督损失的完整梯度。若不使用 stop-gradient, 得到的是 sampled Bellman-residual loss 的梯度;它对应另一个目标,而且对随机 transition 的 Bellman residual 构造无偏梯度通常还会遇到 double-sampling 问题。

k-step TD

实际上,我们可以进一步推广 TD 到多步场景,注意到:

Vπ(s0)=Eπ[r0+γVπ(s1)s0]=Eπ[r0+γEπ[r1+γVπ(s2)s1]s0]=Eπ[r0+γr1+γ2Vπ(s2)s0].\begin{aligned} V^{\pi}(s_0) &= \mathbb{E}^{\pi}[r_0 + \gamma V^{\pi}(s_1) \mid s_0] \\ &= \mathbb{E}^{\pi}[r_0 + \gamma \mathbb{E}^{\pi}[r_1 + \gamma V^{\pi}(s_2) \mid s_1] \mid s_0] \\ &= \mathbb{E}^{\pi}[r_0 + \gamma r_1 + \gamma^2 V^{\pi}(s_2) \mid s_0]. \end{aligned}

因此:

Vπ(s0)=Eπ[r0+γr1+γ2Vπ(s2)s0]V^{\pi}(s_0) = \mathbb{E}^{\pi}[r_0 + \gamma r_1 + \gamma^2 V^{\pi}(s_2) \mid s_0]

重复这个过程 kk 次,我们就可以得到 kk step TD:

Vπ(s0)=Eπ[i=0k1γiri+γkVπ(sk)s0]V^{\pi}(s_0) = \mathbb{E}^{\pi}\left[\sum_{i=0}^{k-1} \gamma^i r_i + \gamma^k V^{\pi}(s_k) \mid s_0\right]

基于 k-step transition ,我们可以构建对应的目标函数并应用 stop-gradient:

Algorithm: Value function approximation with k-step TD

while not converged:

  1. s0p0s_0 \sim p_0, set t=0t = 0.
  2. for t=0,,k1t = 0, \ldots, k-1, stopping early if the episode terminates:
    • atπ(st)a_t \sim \pi(\cdot \mid s_t)
    • (rt,st+1)p(,st,at)(r_t, s_{t+1}) \sim p(\cdot, \cdot \mid s_t, a_t)
  3. Let nkn\leq k be the number of collected transitions and compute y={t=0n1γtrt+γnsg[Vϕ(sn)],if snterm,t=0n1γtrt,otherwise.y = \begin{cases} \sum_{t=0}^{n-1} \gamma^t r_t + \gamma^n \, \mathrm{sg}[V_{\phi}(s_n)], & \text{if } s_n \neq \langle\text{term}\rangle, \\ \sum_{t=0}^{n-1} \gamma^t r_t, & \text{otherwise}. \end{cases}
  4. g(Vϕ(s0)y)ϕVϕ(s0)g \gets (V_{\phi}(s_0) - y) \nabla_{\phi} V_{\phi}(s_0).
  5. Update ϕ\phi using gg with an optimizer.

Q-Function Evaluation

对于 Q-function, 我们也可以和 value function 一样设计类似的算法。对于 MC, 有:

Qπ(s0,a0)=Eπ[t=0T1γtrts0,a0]1Mi=1Mt=0T(i)1γtrt(i).Q^{\pi}(s_0, a_0) =\mathbb{E}^{\pi}\left[\sum_{t=0}^{T-1}\gamma^t r_t\mid s_0,a_0\right] \approx\frac{1}{M}\sum_{i=1}^{M}\sum_{t=0}^{T^{(i)}-1}\gamma^t r_t^{(i)}.

当使用函数来近似时,目标函数为:

L(ϕ)=Esdπ,aπ(s)[12(Qϕ(s,a)Qπ(s,a))2].\mathcal{L}(\phi) = \mathbb{E}_{s \sim d^\pi,\,a\sim\pi(\cdot\mid s)} \left[\frac{1}{2}\left(Q_{\phi}(s,a)-Q^{\pi}(s,a)\right)^2\right].

对应的梯度:

ϕL(ϕ)=Eπ[(Qϕ(st,at)Gt)ϕQϕ(st,at)].\nabla_{\phi} \mathcal{L}(\phi) =\mathbb{E}^{\pi}\left[ \left(Q_{\phi}(s_t,a_t)-G_t\right)\nabla_{\phi}Q_{\phi}(s_t,a_t) \right].

对应的算法如下

Algorithm: Q-function approximation with MC

while not converged:

  1. s0p0s_0 \sim p_0, set t=0t = 0, τ=[]\tau=[].
  2. while stterms_t \neq \langle\text{term}\rangle:
    • atπ(st)a_t \sim \pi(\cdot \mid s_t)
    • (rt,st+1)p(,st,at)(r_t, s_{t+1}) \sim p(\cdot, \cdot \mid s_t, a_t)
    • τ.append((st,at,rt))\tau.\mathrm{append}((s_t,a_t,r_t))
    • tt+1t \gets t + 1
  3. Set T=tT = t.
  4. For every tt, compute Gt=k=tT1γktrkG_t=\sum_{k=t}^{T-1}\gamma^{k-t}r_k.
  5. g=1Tt=0T1(Qϕ(st,at)Gt)ϕQϕ(st,at)g=\frac{1}{T}\sum_{t=0}^{T-1} (Q_\phi(s_t,a_t)-G_t)\nabla_\phi Q_\phi(s_t,a_t).
  6. Update ϕ\phi using gg with an optimizer.

Return QϕQ_\phi

对于 TD learning, 我们的目标函数为

L(ϕ)=Esdπ,aπ(s)[12(Qϕ(s,a)rγQπ(s,a))2],\mathcal{L}(\phi) = \mathbb{E}_{s \sim d^\pi,\,a\sim\pi(\cdot\mid s)} \left[\frac{1}{2}\left(Q_{\phi}(s,a)-r-\gamma Q^\pi(s',a')\right)^2\right],

其中非终止状态下 aπ(s)a'\sim\pi(\cdot\mid s'). 与 value function 类似,使用 bootstrap target 可以得到 on-policy 的 Expected SARSA/SARSA 风格 policy evaluation:

Algorithm: Q-function approximation with TD

While not converged:

  1. s0p0s_0 \sim p_0, a0π(s0)a_0 \sim \pi(\cdot \mid s_0), (r0,s1)p(,s0,a0)(r_0, s_1) \sim p(\cdot, \cdot \mid s_0, a_0).
  2. y={r0+γsg[Qϕ(s1,a1)],a1π(s1),if s1term,r0,otherwise.y = \begin{cases} r_0 + \gamma \, \mathrm{sg}[Q_{\phi}(s_1, a_1)], \quad a_1\sim\pi(\cdot\mid s_1), & \text{if } s_1 \neq \langle\text{term}\rangle, \\ r_0, & \text{otherwise}. \end{cases}
  3. g=(Qϕ(s0,a0)y)ϕQϕ(s0,a0)g = \left(Q_{\phi}(s_0, a_0) - y\right) \nabla_{\phi} Q_{\phi}(s_0, a_0).
  4. Update ϕ\phi using gg with an optimizer.

Return QϕQ_\phi

Algorithm: Q-function approximation with k-step TD

While not converged:

  1. s0p0s_0 \sim p_0, set t=0t = 0.
  2. For t=0,,k1t = 0, \ldots, k-1, stopping early on termination:
    • atπ(st)a_t \sim \pi(\cdot \mid s_t)
    • (rt,st+1)p(,st,at)(r_t, s_{t+1}) \sim p(\cdot, \cdot \mid s_t, a_t)
  3. Let nkn\leq k be the number of collected transitions.
  4. If sns_n is non-terminal
    1. sample anπ(sn)a_n\sim\pi(\cdot\mid s_n)
    2. set y=t=0n1γtrt+γnsg[Qϕ(sn,an)]y=\sum_{t=0}^{n-1}\gamma^t r_t+\gamma^n\mathrm{sg}[Q_\phi(s_n,a_n)];
  5. Else
    1. set y=t=0n1γtrty=\sum_{t=0}^{n-1}\gamma^t r_t.
  6. g=(Qϕ(s0,a0)y)ϕQϕ(s0,a0)g = (Q_{\phi}(s_0, a_0) - y) \nabla_{\phi} Q_{\phi}(s_0, a_0)
  7. Update ϕ\phi using gg with an optimizer.

Return QϕQ_\phi

Policy Iteration Methods

介绍完 policy evaluation 之后,我们可以设计 policy iteration 算法,其基本思想是交替进行 policy evaluation 和 policy improvement:

Compute Vπk and Qπk(Policy evaluation)πk+1(s)=argmaxaAE(r,s)p(,s,a)[r+γVπk(s)s,a](Policy improvement)\begin{aligned} &\text{Compute } V^{\pi^k} \text{ and } Q^{\pi^k} &\quad&\text{(Policy evaluation)} \\ &\pi^{k+1}(s) = \arg\max_{a \in \mathcal{A}} \mathbb{E}_{(r, s') \sim p(\cdot, \cdot \mid s, a)}\left[r + \gamma V^{\pi^k}(s') \mid s, a\right] &\quad&\text{(Policy improvement)} \end{aligned}

算法的正确性由如下定理保证

Theorem:

Theorem: 考虑上述 policy iteration 算法,假设 γ(0,1)\gamma \in (0, 1), S<|\mathcal{S}| < \infty, A<|\mathcal{A}| < \infty 以及 rR<|r| \leq R < \infty, a.s. 那么:

Vπk+1Vπk,k=1,2,V^{\pi^{k+1}} \geq V^{\pi^k}, \quad k = 1, 2, \dots

并且存在 KNK \in \mathbb{N} 满足 VπK=VV^{\pi^K} = V^*. 也就是说,有限步之后就可以得到最优策略.

Proof:

πk+1\pi^{k+1} 的 greedy 定义,对每个状态 ss 都有

(Tπk+1Vπk)(s)=(TVπk)(s)(TπkVπk)(s)=Vπk(s).(\mathcal T^{\pi^{k+1}}V^{\pi^k})(s) =(\mathcal T^*V^{\pi^k})(s) \geq(\mathcal T^{\pi^k}V^{\pi^k})(s) =V^{\pi^k}(s).

反复应用单调的 Bellman operator Tπk+1\mathcal T^{\pi^{k+1}},并利用其 contraction 性质,得到

Vπk+1=limn(Tπk+1)nVπkVπk.V^{\pi^{k+1}} =\lim_{n\to\infty} (\mathcal T^{\pi^{k+1}})^nV^{\pi^k} \geq V^{\pi^k}.

有限 MDP 只有有限个 deterministic policies. 若 value 尚未最优,policy improvement 会产生严格改进;采用一致的 tie-breaking 后不会在等价 policy之间循环。 因此有限次迭代后到达满足 Bellman optimality equation 的 policy, 其 value 即为 VV^*.

MC v.s. TD

MC 必须等待 episode 结束,然后用真实 return 监督每个已访问状态。 对固定状态而言,sample return 是 VπV^\pi 的无偏估计,但长轨迹中的 variance 可能很高。

TD 使用 Vϕ(st+1)V_\phi(s_{t+1}) bootstrap,不需要等待完整 episode,通常 variance 更低,但 target 会受到当前 value approximation error 的影响。 这里所谓的 “TD 有偏” 主要是指有限训练阶段的 bootstrap target;在满足收敛条件的 tabular on-policy 场景中,TD 仍然可以收敛到真实的 VπV^\pi.

k-step TD 在两者之间插值:

TD(0)k0k-step TDkTMC.\text{TD(0)} \xleftarrow{\quad k\to0\quad} \text{k-step TD} \xrightarrow{\quad k\to T\quad} \text{MC}.

Takeaway

Value Based Methods

Policy Evaluation 一节中,策略 π\pi 是固定的,我们只估计 VπV^\piQπQ^\pi. Value-based methods 更进一步:学习能够区分动作优劣的 value function,并据此改进策略。

一种直接思路是寻找 optimal value function VV^* 或 optimal Q function QQ^*,再从中恢复最优策略。 这也是 value based methods 的核心思想:

求解得到 optimal value function VV^* 或者 optimal Q function QQ^*, 再得到 optimal policy π\pi^*.

Value Iteration

Value iteration 的基本思路为:通过 Bellman optimality operator 迭代求解 VV^*QQ^*, 然后得到最优策略。 根据 BOE Theorem 的证明,T\mathcal{T}^* 是一个 contraction mapping, 因此我们可以构造函数序列:

V,TV,(T)2V,V, \mathcal{T}^* V, (\mathcal{T}^*)^2 V, \cdots

由不动点定理,该序列收敛到 VV^*. 这就是 value iteration algorithm:

Vk+1=TVkVk+1(st)=maxaAE(rt,st+1)p(,st,at)[r+γVk(st+1)St=st,At=at]\begin{aligned} V^{k+1} &= \mathcal{T}^* V^k \\ V^{k+1}(s_t) &= \max_{a \in \mathcal{A}} \mathbb{E}_{(r_t, s_{t+1}) \sim p(\cdot, \cdot \mid s_t, a_t)}\left[r + \gamma V^k(s_{t+1}) \mid S_t=s_t, A_t=a_t\right] \end{aligned}

类似地,我们还有 Q value iteration algorithm:

Qk+1=TQkQk+1(st,at)=E(rt,st+1)p(,st,at)[r+γmaxat+1AQk(st+1,at+1)St=st,At=at]\begin{aligned} Q^{k+1} &= \mathcal{T}^* Q^k \\ Q^{k+1}(s_t, a_t) &= \mathbb{E}_{(r_t, s_{t+1}) \sim p(\cdot, \cdot \mid s_t, a_t)}\left[r + \gamma \max_{a_{t+1} \in \mathcal{A}} Q^k(s_{t+1}, a_{t+1}) \mid S_t=s_t,A_t=a_t\right] \end{aligned}

最终,我们的最优策略为

π(st)=argmaxaAQ(st,a)\pi^{*}(s_t) = \arg\max_{a \in \mathcal{A}} Q^*(s_t,a)

其中 Q=limkQkQ^*=\lim_{k\to\infty}Q^k.

本节将 value iteration 的思想与函数逼近、TD learning 等技术结合,得到实用的 value-based 算法.

Tabular Learning

当 state 和 action space 足够小时,我们可以直接维护一张 Q-table.

在 Model-Free 设定下,我们无法直接计算转移概率 Esp\mathbb{E}_{s' \sim p}. 此时,我们必须通过与环境交互产生的轨迹样本 (s,a,r,s,a)(s, a, r, s', a') 来更新价值。

SARSA

SARSA 的核心在于它的更新动作 aa' 严格来自于当前的 ϵ\epsilon-greedy 策略:

Q(s,a)Q(s,a)+α[r+γQ(s,a)Q(s,a)]\boxed{Q(s, a) \leftarrow Q(s, a) + \alpha \left[ r + \gamma Q(s', a') - Q(s, a) \right]}

Q-Learning (Off-Policy TD Control)

与 SARSA 不同,Q-Learning 在估计未来价值时,直接贪心地选择最优动作,而不受当前探索策略的影响:

Q(s,a)Q(s,a)+α[r+γmaxaAQ(s,a)Q(s,a)]\boxed{Q(s, a) \leftarrow Q(s, a) + \alpha \left[ r + \gamma \max_{a' \in \mathcal{A}} Q(s', a') - Q(s, a) \right]}

Function Approximation

当状态空间较大或连续时,我们需要使用函数逼近,与上一节类似,我们用 VϕV_{\phi} 近似 VπV^{\pi}:

J(ϕ)=Es[12(Vπ(s)Vϕ(s))2]J(\phi) = \mathbb{E}_{s}\left[\frac{1}{2}\left(V^{\pi}(s) - V_{\phi}(s)\right)^2\right]

目标函数的梯度为:

ϕJ(ϕ)=Es[(Vπ(s)Vϕ(s))ϕVϕ(s)]\nabla_{\phi} J(\phi) = -\mathbb{E}_{s}\left[(V^{\pi}(s) - V_{\phi}(s))\nabla_{\phi} V_{\phi}(s)\right]

随机梯度下降的更新形式为:

ϕk+1=ϕk+αk(Vπ(s)Vϕ(s))ϕVϕ(s)\phi_{k+1} = \phi_k + \alpha_k (V^{\pi}(s) - V_{\phi}(s))\nabla_{\phi} V_{\phi}(s)

Vπ(s)V^{\pi}(s) 可以通过两种方式来近似:

  1. MC 近似: Vπ(s)=Eπ[G0s0=s]G0V^{\pi}(s) = \mathbb{E}^{\pi}[G_0 \mid s_0 = s] \approx G_0
  2. TD 近似: Vπ(s)rt+γVϕ(st+1)V^{\pi}(s) \approx r_t + \gamma V_{\phi}(s_{t+1})

Deep Q-Network (DQN)

将 function approximation 与 Q-learning 结合,可以用 neural network 逼近 optimal Q function。DQN 进一步引入 replay buffer 和 target network 来改善训练稳定性:

L(ϕ)=E(s,a,r,s)D[12(r+γmaxaQϕ(s,a)Qϕ(s,a))2]\mathcal{L}(\phi) = \mathbb{E}_{(s,a,r,s') \sim \mathcal{D}} \left[ \frac{1}{2} \left( r + \gamma \max_{a'} Q_{\phi^-}(s', a') - Q_{\phi}(s, a) \right)^2 \right]

其中 ϕ\phi^- 为**目标网络(Target Network)**的参数,用于切断数据自举(Bootstrapping)带来的训练正反馈爆炸。 为了保证求导时 Target 的稳定性,我们必须在代码或公式里明确使用 stop-gradient 算子:

y=r+γmaxasg[Qϕ(s,a)]y = r + \gamma \max_{a'} \mathrm{sg}\left[Q_{\phi^-}(s', a')\right]

DQN 针对 optimal policy 应用了 TD learning, 其目标函数为:

minϕL(ϕ)=E(s,a,r,s)D[12(Qϕ(s,a)(r+γmaxaQϕ(s,a)))2]\min_{\phi} \mathcal{L}(\phi) = \mathbb{E}_{(s, a, r, s') \sim \mathcal{D}}\left[\frac{1}{2}\left(Q_{\phi}(s, a) - \left(r + \gamma \max_{a'} Q_{\phi}(s', a')\right)\right)^2\right]

这里 D\mathcal{D} 是 replay buffer, 存储了过去的 transition (s,a,r,s)(s, a, r, s').

DQN 的两个关键技术:

  1. Experience Replay: 使用 replay buffer 存储过去的 transitions,从中随机采样进行训练。这打破了数据之间的相关性,提高了样本效率.
  2. Target Network: 使用一个独立的 target network QϕQ_{\phi^{-}} 来计算 TD target, 其参数 ϕ\phi^{-} 定期从 ϕ\phi 复制(或通过 Polyak averaging 更新)。这稳定了训练过程.

带有 target network 和 stop-gradient 的 DQN 目标函数为:

L(ϕ)=E(s,a,r,s)D[12(Qϕ(s,a)(r+γmaxasg[Qϕ(s,a)]))2]\mathcal{L}(\phi) = \mathbb{E}_{(s, a, r, s') \sim \mathcal{D}}\left[\frac{1}{2}\left(Q_{\phi}(s, a) - \left(r + \gamma \, \max_{a'} \mathrm{sg}[Q_{\phi^{-}}(s', a')]\right)\right)^2\right]

对应的 Python 实现:

# Compute current Q-values
q_values = Q_phi(states).gather(1, actions)

# Compute target Q-values with target network
with torch.no_grad():
    next_q_values = Q_target(next_states).max(1)[0]
    targets = rewards + gamma * next_q_values * (1 - dones)

# Compute loss
loss = 0.5 * (q_values - targets) ** 2
loss = loss.mean()
Algorithm: DQN with Experience Replay
  1. Initialize replay buffer D\mathcal{D} with capacity NN.
  2. Initialize Q-network QϕQ_{\phi} with random weights.
  3. Initialize target network QϕQ_{\phi^{-}} with weights ϕ=ϕ\phi^{-} = \phi.
  4. for episode = 1,,M1, \dots, M:
    • s0p0s_0 \sim p_0.
    • for t=0,,T1t = 0, \dots, T-1:
      • Select action with ϵ\epsilon-greedy: at={argmaxaQϕ(st,a),with prob 1ϵrandom action,with prob ϵa_t = \begin{cases} \arg\max_a Q_{\phi}(s_t, a), & \text{with prob } 1 - \epsilon \\ \text{random action}, & \text{with prob } \epsilon \end{cases}
      • (rt,st+1)p(,st,at)(r_t, s_{t+1}) \sim p(\cdot, \cdot \mid s_t, a_t).
      • Store (st,at,rt,st+1)(s_t, a_t, r_t, s_{t+1}) in D\mathcal{D}.
      • Sample random minibatch {(si,ai,ri,si)}\{(s_i, a_i, r_i, s'_i)\} from D\mathcal{D}.
      • Compute target: yi={ri+γmaxaQϕ(si,a),if not terminalri,if terminaly_i = \begin{cases} r_i + \gamma \max_{a'} Q_{\phi^{-}}(s'_i, a'), & \text{if not terminal} \\ r_i, & \text{if terminal} \end{cases}
      • Update Q-network: L=1batchi12(Qϕ(si,ai)yi)2,ϕϕαϕL\mathcal{L} = \frac{1}{|\text{batch}|} \sum_i \frac{1}{2} (Q_{\phi}(s_i, a_i) - y_i)^2, \quad \phi \gets \phi - \alpha \nabla_{\phi} \mathcal{L}
      • Every CC steps: ϕϕ\phi^{-} \gets \phi.

Why LLMs Usually Do Not Use DQN

Value-based control 建立了从 Bellman optimality equation 到可训练算法的桥梁, 但 DQN 通常不是 LLM post-training 的最终选择:

  1. LLM action space 是整个 vocabulary,每个 prefix 都要对大量 token 计算并比较 Q(s,a)Q(s,a)
  2. argmaxaQ(s,a)\arg\max_a Q(s,a) 产生的 greedy policy 不等于我们希望保留和微调的 autoregressive policy distribution
  3. LLM reward 常在完整 response 后给出,早期 token 的 Q target 具有严重的 credit-assignment 和 exploration 问题
  4. 预训练 LLM 已经提供了强 policy prior,直接优化 πθ\pi_\theta 通常比从 Q function 重新构造生成策略更自然

因此,本章的主要作用是解释 evaluation 如何变成 control, 并为 TD, off-policy learning 和 bootstrap 建立直觉。 下一章将转向 policy-based methods:直接对 LLM policy 的 expected return 求梯度。

Policy Based Methods

Value-based methods 先学习 QQ^*,再通过 argmaxaQ(s,a)\arg\max_a Q^*(s,a) 间接得到策略。 对 LLM 而言,我们已经有一个参数化的 autoregressive policy πθ(ytx,y<t)\pi_\theta(y_t\mid x,y_{<t}), 因此更自然的选择是直接调整 θ\theta,提高高回报 response 的生成概率。 这就是 policy-based methods 的核心思想。

Core idea

使用 policy πθ\pi_\theta 生成 trajectory, 通过 trajectory reward 构造 gradient estimator,再直接改进当前策略 πθ\pi_\theta.

根据采样策略 (behavior policy) 的不同,policy-based methods 一般会被分为两类:

  1. on-policy: 采集数据的 policy 和训练更新的 policy 一致
  2. off-policy: 采集数据的 policy 和训练更新的 policy 不一致

这两者的核心区别在于:我们更新模型所使用的数据是否由当前模型产生? 如果是,则说明算法是 On-policy, 反之则说明是 Off-policy.

on-policy 和 off-policy 对比如下表所示

Categoryon-policyoff-policy
examplesPPO, GRPODQN, DDPG
sampled fromcurrent policybehavior policy
data explorationcontrolled by current policycontrolled by behavior policy
infra efficiencylow (needs to generate new sample)high
training stabilitysensitive to policy update sizesensitive to distribution shift
compute costhighlow
best forfresh rollout and direct policy optimizationreusing past or external data

目前,主流针对 LLM 的 RL 算法基本都是

Objective

我们的目标与前面一致,仍然是最大化 return, 问题形式为

maxπJ(θ)=Eτπθ[G0(τ)]\max_{\pi}\mathcal{J}(\theta)=\mathbb{E}_{\tau\sim\pi_\theta}\left[G_0(\tau)\right]

Policy Gradient Theorem

我们先来求解目标函数的梯度,梯度由下面的 policy gradient theorem 给出。

Theorem: Policy Gradient Theorem

目标函数 最大化 return 的梯度为

θJ(θ)=Eτπθ[t=0T1θlogπθ(atst)G0(τ)]\nabla_\theta \mathcal{J}(\theta) = \mathbb{E}_{\tau\sim\pi_\theta}\left[ \sum_{t=0}^{T-1}\nabla_\theta\log\pi_\theta(a_t\mid s_t)G_0(\tau)\right]
Proof:

trajectory definition, trajectory τ\tau 的概率分布可以写为

pθ(τ)=p0(s0)t=0T1p(rt,st+1st,at)πθ(atst)p_\theta(\tau) = p_0(s_0)\prod_{t=0}^{T-1} p(r_t,s_{t+1}\mid s_t,a_t)\pi_\theta(a_t\mid s_t)

pθ(τ)p_\theta(\tau) 求梯度得到

θpθ(τ)=pθ(τ)θlogpθ(τ)=pθ(τ)θlog[p0(s0)t=0T1p(rt,st+1st,at)πθ(atst)]=pθ(τ)t=0T1logπθ(atst)\begin{aligned} \nabla_\theta p_\theta(\tau)&=p_\theta(\tau)\nabla_\theta \log p_\theta(\tau)\\ &= p_\theta(\tau)\nabla_\theta \log \left[p_0(s_0)\prod_{t=0}^{T-1} p(r_t,s_{t+1}\mid s_t,a_t)\pi_\theta(a_t\mid s_t)\right]\\ &= p_\theta(\tau)\sum_{t=0}^{T-1}\log\pi_\theta(a_t\mid s_t) \end{aligned}

这样,目标函数的梯度为

θJ(θ)=θEτπθ[G0(τ)]=τG0(τ)θpθ(τ)=τG0(τ)pθ(τ)t=0T1logπθ(atst)=Eτπθ[t=0T1θlogπθ(atst)G0(τ)]\begin{aligned} \nabla_\theta \mathcal{J}(\theta)&=\nabla_\theta\mathbb{E}_{\tau\sim\pi_\theta}\left[G_0(\tau)\right]\\ &= \sum_\tau G_0(\tau)\nabla_\theta p_\theta(\tau)\\ &= \sum_\tau G_0(\tau)p_\theta(\tau)\sum_{t=0}^{T-1}\log\pi_\theta(a_t\mid s_t) \\ &= \mathbb{E}_{\tau\sim\pi_\theta}\left[ \sum_{t=0}^{T-1}\nabla_\theta\log\pi_\theta(a_t\mid s_t)G_0(\tau)\right] \end{aligned}

证毕。

注意 policy gradient theorem 将梯度也写成了期望的形式,这方面我们使用 MC 方法来进行估计和优化。

Vanilla Policy Gradient Methods

接下来,我们介绍最简单的 policy gradient methods, 令

g^:=t=0T1θlogπθ(atst)G0(τ)\hat{g} := \sum_{t=0}^{T-1}\nabla_\theta\log\pi_\theta(a_t\mid s_t)G_0(\tau)

由 policy gradient theorem, 我们知道 g^\hat{g}θJ(θ)\nabla_\theta \mathcal{J}(\theta) 的一个无偏估计。

基于 MC 思想,我们可以得到最简单的 policy gradient algorithm, 也就是 REINFORCE 算法

Algorithm: Policy Gradient with MC (REINFORCE)

Initialize: policy parameters θ0\theta_0.

For k=0,1,k=0,1,\dots:

  • Sample MM trajectories {τi}\{\tau_i\} according to policy πθk\pi_{\theta_k}.
  • Estimate the gradient via MC: g^=1Mi=1Mt=0Ti1θlogπθk(atst)G0(τi)\hat{g}=\frac{1}{M}\sum_{i=1}^M\sum_{t=0}^{T_i-1}\nabla_\theta\log\pi_{\theta_k}(a_t\mid s_t)G_0(\tau_i)
  • Update θk\theta_k using g^\hat{g} with an optimizer.

Return πθk\pi_{\theta_k}

Variance Reduction

虽然 ggθJ(θ)\nabla_\theta \mathcal{J}(\theta) 的一个无偏估计,但是使用 gg 作为估计会出现比较大的方差,TODO

因此我们本节主要介绍如何减少 gg 作为估计的方差,即 variance reduction.

我们首先介绍 policy gradient 的 baseline invariance 性质。

Baseline Invariance

Proposition: Baseline Invariance of Policy Gradient

针对 policy gradient theorem, 我们有如下等式

Eτ(p0,πθ,p)[t=0T1θlogπθ(atst)t=0T1r(st)]=Eτ(p0,πθ,p)[t=0T1θlogπθ(atst)(t=0T1r(st)b(st))]\begin{aligned} &\mathbb{E}_{\tau\sim(p_0,\pi_\theta, p)}\left[ \sum_{t=0}^{T-1}\nabla_\theta\log\pi_\theta(a_t\mid s_t)\sum_{t=0}^{T-1}r(s_t)\right] \\ =& \mathbb{E}_{\tau\sim(p_0,\pi_\theta, p)}\left[ \sum_{t=0}^{T-1}\nabla_\theta\log\pi_\theta(a_t\mid s_t)\left(\sum_{t=0}^{T-1}r(s_t)-b(s_t)\right)\right] \end{aligned}

其中 b(st)b(s_t) 是一个仅与状态有关的函数或者常数,我们将其称为 baseline.

Proof:

我们证明

Eτ(p0,πθ,p)[b(st)θlogpθ(atst)]=0\mathbb{E}_{\tau\sim(p_0,\pi_\theta, p)}\left[b(s_t)\nabla_\theta \log p_\theta(a_t\mid s_t)\right]=0

注意到

Eτ(p0,πθ,p)[b(st)θlogpθ(atst)]=Es0p0[b(st)Eatπθ(st)[θlogπθ(atst)]]=Es0p0[b(st)aθπθ(as)]=Es0p0[b(st)θaπθ(as)]=Es0p0[b(st)θ1]=0\begin{aligned} \mathbb{E}_{\tau\sim(p_0,\pi_\theta, p)}\left[b(s_t)\nabla_\theta \log p_\theta(a_t\mid s_t)\right] &= \mathbb{E}_{s_0\sim p_0}\left[b(s_t)\mathbb{E}_{a_t\sim\pi_\theta(\cdot\mid s_t)}\left[\nabla_\theta \log \pi_\theta(a_t\mid s_t)\right]\right]\\ &= \mathbb{E}_{s_0\sim p_0}\left[b(s_t)\sum_{a}\nabla_\theta \pi_\theta(a\mid s)\right]\\ &= \mathbb{E}_{s_0\sim p_0}\left[b(s_t)\nabla_\theta \sum_{a}\pi_\theta(a\mid s)\right]\\ &= \mathbb{E}_{s_0\sim p_0}\left[b(s_t)\nabla_\theta 1\right]\\ &=0 \end{aligned}

因此,policy gradient 满足 baseline invariance property.

另一方面,Rao–Blackwell Theorem 说明了我们可以通过构造一个新的无偏估计,这个新的无偏估计相对于原始的无偏估计,有更小的方差。

Theorem: Rao-Blackwell Theorem

XXYY 为随机变量, 令 I^1(X,Y)\hat{I}_1(X,Y)II 的一个无偏估计,即

I=EXY[I^1(XY)]I = \mathbb{E}_{XY}[\hat{I}_1(XY)]

I^2(Y)=EXY[I^1(X,Y)Y]\hat{I}_2(Y)=\mathbb{E}_{X\mid Y}[\hat{I}_1(X,Y)\mid Y], 则 I^2\hat{I}_2 也是 II 的一个无偏估计,并且

var[I^2]var[I^1]\mathrm{var}[\hat{I}_2]\leq\mathrm{var}[\hat{I}_1]

I^2\hat{I}_2 被称为 I^1(X,Y)\hat{I}_1(X,Y)Rao-Blackwellized estimator.

Proof:

Expectation 中的 law of total expectation, 我们有

EXY[I^2(XY)]=EY[EXY[I^1(X,Y)Y]]=EXY[I^1(XY)]=I\mathbb{E}_{XY}[\hat{I}_2(XY) ]=\mathbb{E}_{Y}[\mathbb{E}_{X\mid Y}[\hat{I}_1(X,Y)\mid Y]]=\mathbb{E}_{XY}[\hat{I}_1(XY) ]=I

其次,我们有

var[I^2]=EY[(II^2(Y))2]=EY[(EXY[II^1(X,Y)Y])2]EY[(EXY[(II^1(X,Y))2Y])]=EY[(II^1(X,Y))2]=var[I^1]\begin{align} \mathrm{var}[\hat{I}_2]&=\mathbb{E}_{Y}[(I-\hat{I}_2(Y))^2]\\ &= \mathbb{E}_{Y}[\left(\mathbb{E}_{X\mid Y}\left[I-\hat{I}_1(X,Y)\mid Y\right]\right)^2]\\ &\leq\mathbb{E}_{Y}[\left(\mathbb{E}_{X\mid Y}\left[(I-\hat{I}_1(X,Y))^2\mid Y\right]\right)]\\ &= \mathbb{E}_{Y}[(I-\hat{I}_1(X,Y))^2]\\ &=\mathrm{var}[\hat{I}_1] \end{align}

这里不等式使用了 Jensen’s inequality. 证毕。\blacksquare

基于 Rao–Blackwell Theorem, 我们可以通过改变 b(st)b(s_t) 来调整 policy gradient 的 variance.

Reward to Go

Reward to go 的基本思想为

Core idea

对于时刻 tt, 过去的轨迹 s0,r0,,rt1s_0,r_0,\dots,r_{t - 1} 不应该被考虑,因为它们都已经固定了,加入它们会产生额外的 variance, 此时,我们对应的 b(st)b(s_t) 定义为

其定义的 baseline 如下所示

b(st)=k=0tγkrkb(s_t) = \sum_{k=0}^{t}\gamma^k r_k

此时,我们的 policy gradient 为

θJ(θ)=Eτ(p0,πθ,p)[t=0T1θlogπθ(atst)k=tT1r(sk)]\nabla_\theta \mathcal{J}(\theta)=\mathbb{E}_{\tau\sim(p_0,\pi_\theta, p)}\left[ \sum_{t=0}^{T-1}\nabla_\theta\log\pi_\theta(a_t\mid s_t)\sum_{k=t}^{T-1}r(s_{k})\right]

对应的 REINFORCE 算法改进版本为

Algorithm: Policy Gradient with MC and reward to go (REINFORCE)
  1. Initialize policy parameters θ0\theta_0.
  2. for k=0,1,k=0,1,\dots:
    • Sample MM trajectories {τi}\{\tau_i\} according to policy πθk\pi_{\theta_k}.
    • Estimate the gradient via MC: g^=1Mi=1Mθlogπθk(atst)γtGt(τi)\hat{g}=\frac{1}{M}\sum_{i=1}^M\nabla_\theta\log\pi_{\theta_k}(a_t\mid s_t)\gamma^tG_t(\tau_i)
    • Update θk\theta_k using g^\hat{g} with an optimizer.

Optimal Baseline

我们可以从理论上推导出最优 baseline.

Lemma:

ssaa 为随机变量,w(s,a)w(s,a), Q(s,a)Q(s,a), b(s)b(s) 为函数,则

b=argminb()Ea,s[w2(s,a)(Q(s,a)b(s))2]b^* = \arg\min_{b(\cdot)} \mathbb{E}_{a,s}\left[w^2(s,a)\left(Q(s,a)-b(s)\right)^2\right]

其中

b(s)=Eas[w2(s,a)Q(s,a)s]Eas[w2(s,a)s]b^*(s) = \frac{\mathbb{E}_{a\mid s}\left[w^2(s,a)Q(s,a)\mid s\right]}{\mathbb{E}_{a\mid s}\left[w^2(s,a)\mid s\right]}
Proof:

对目标函数进行展开得到

Ea,s[w2(s,a)(Q(s,a)b(s))2]=Ea,s[w2(s,a)(Q(s,a)b(s)+b(s)b(s))2]=Ea,s[w2(s,a)(Q(s,a)b(s))2]+Ea,s[w2(s,a)(b(s)b(s))2]+2Ea,s[w2(s,a)(Q(s,a)b(s))(b(s)b(s))]=Ea,s[w2(s,a)(Q(s,a)b(s))2]+Ea,s[w2(s,a)(b(s)b(s))2]Ea,s[w2(s,a)(Q(s,a)b(s))2]\begin{aligned} \mathbb{E}_{a,s}\left[w^2(s,a)\left(Q(s,a)-b(s)\right)^2\right]&= \mathbb{E}_{a,s}\left[w^2(s,a)\left(Q(s,a)-b^*(s)+b^*(s)-b(s)\right)^2\right]\\ &= \mathbb{E}_{a,s}\left[w^2(s,a)\left(Q(s,a)-b^*(s)\right)^2\right]+\mathbb{E}_{a,s}\left[w^2(s,a)\left(b^*(s)-b(s)\right)^2\right]\\ &+2\mathbb{E}_{a,s}\left[w^2(s,a)\left(Q(s,a)-b^*(s)\right)\left(b^*(s)-b(s)\right)\right]\\ &= \mathbb{E}_{a,s}\left[w^2(s,a)\left(Q(s,a)-b^*(s)\right)^2\right]+\mathbb{E}_{a,s}\left[w^2(s,a)\left(b^*(s)-b(s)\right)^2\right]\\ &\geq \mathbb{E}_{a,s}\left[w^2(s,a)\left(Q(s,a)-b^*(s)\right)^2\right] \end{aligned}

其中第三个等式用到了

Ea,s[w2(s,a)(Q(s,a)b(s))(b(s)b(s))]=Es[Eas[w2(s,a)(Q(s,a)b(s))(b(s)b(s))]s]=Es[(Eas[w2(s,a)Q(s,a)s]Eas[w2(s,a)s]b(s))(b(s)b(s))]=0\begin{aligned} \mathbb{E}_{a,s}\left[w^2(s,a)\left(Q(s,a)-b^*(s)\right)\left(b^*(s)-b(s)\right)\right]&= \mathbb{E}_{s}\left[\mathbb{E}_{a\mid s}\left[w^2(s,a)\left(Q(s,a)-b^*(s)\right)\left(b^*(s)-b(s)\right)\right]\mid s\right]\\ &= \mathbb{E}_{s}\left[\left(\mathbb{E}_{a\mid s}\left[w^2(s,a)Q(s,a)\mid s\right] - \mathbb{E}_{a\mid s}\left[w^2(s,a)\mid s\right]b^*(s)\right)\left(b^*(s)-b(s)\right)\right]\\ &=0 \end{aligned}

证毕。\blacksquare

基于上述引理,我们可以得到关于 baseline invariance 中的最优 baseline:

b(s)=Eatπθ(st)[(θlogπθ(atst))2Qπθ(s,a)s]Eatπθ(st)[θlogπθ(atst))2s]b^*(s) =\frac{\mathbb{E}_{a_t\sim\pi_\theta(\cdot\mid s_t)}\left[(\nabla_\theta\log\pi_\theta(a_t\mid s_t))^2Q^{\pi_\theta}(s,a)\mid s\right]}{\mathbb{E}_{a_t\sim\pi_\theta(\cdot\mid s_t)}\left[\nabla_\theta\log\pi_\theta(a_t\mid s_t))^2\mid s\right]}

Advantage

理论上我们可以使用 b(st)b^*(s_t) 作为 baseline,但是实际上由于其表达式太复杂,我们一般使用其简化版本。 简化版本移除了 b(s)b^*(s) 中的 θlogπθ(atst)\nabla_\theta\log\pi_\theta(a_t\mid s_t), 这样我们的估计就是

b(s)=Eatπθ(st)[Qπθ(s,a)s]Eatπθ(st)[1s]=Vπθ(s)b(s) = \frac{\mathbb{E}_{a_t\sim\pi_\theta(\cdot\mid s_t)}\left[Q^{\pi_\theta}(s,a)\mid s\right]}{\mathbb{E}_{a_t\sim\pi_\theta(\cdot\mid s_t)}\left[1\mid s\right]} = V^{\pi_\theta}(s)

b(s)=Vπθ(s)b(s)=V^{\pi_\theta}(s) 并不是最优 baseline, 但是其性质比较好。

我们将 b(s)=Vπθ(s)b(s)=V^{\pi_\theta}(s) 带入到目标函数中就得到经典的 policy gradient 表达式

θJ(θ)=Eτ(p0,πθ,p)[t=0T1θlogπθ(atst)(Qπθ(st,at)Vπθ(st))]\nabla_\theta \mathcal{J}(\theta)=\mathbb{E}_{\tau\sim(p_0,\pi_\theta, p)}\left[ \sum_{t=0}^{T-1}\nabla_\theta\log\pi_\theta(a_t\mid s_t)\left(Q^{\pi_\theta}(s_t,a_t)-V^{\pi_\theta}(s_t)\right)\right]

我们定义

Aπθ(st,at):=Qπθ(st,at)Vπθ(st)\boxed{ A^{\pi_\theta}(s_t,a_t) := Q^{\pi_\theta}(s_t,a_t)-V^{\pi_\theta}(s_t) }

ata_tsts_t 处的 advantage. 注意到 Vπθ(st)V^{\pi_\theta}(s_t) 评估了当前状态下沿着策略 πθ\pi_\theta 所能够获得的 discounted return, 因此 advantage 说明了 action ata_t 对于 πθ\pi_\theta 的平均好坏程度:

实际场景下,由于 Vπθ(st)V^{\pi_\theta}(s_t) 未知,使用前面的 value iteration 方法,我们可以使用估计

A^πθ(st,at):=Qπθ(st,at)Vϕ(st)\hat{A}^{\pi_\theta}(s_t,a_t) := Q^{\pi_\theta}(s_t,a_t)-V^{\phi}(s_t)

来代替。

对于最优策略 π\pi^*,我们有

Aπ(st,at):=Qπ(st,at)Vπ(st)0A^{\pi^*}(s_t,a_t) := Q^{\pi^*}(s_t,a_t)-V^{\pi^*}(s_t)\leq 0

证明也很简单,注意到 Vπ(st)=maxaQπ(st,at)V^{\pi^*}(s_t)=\max_{a}Q^{\pi^*}(s_t,a_t) 即可。

最后,我们的 REINFORCE 算法改进如下

Algorithm: Policy Gradient with MC and value function (REINFORCE)
  1. Initialize policy parameters θ0\theta_0.
  2. for k=0,1,k=0,1,\dots:
    • Sample MM trajectories {τi}\{\tau_i\} according to policy πθk\pi_{\theta_k}.
    • Estimate the gradient via MC: g^=1Mi=1Mθlogπθk(atst)γt(Qπθ(st,at)Vϕ(st))\hat{g}=\frac{1}{M}\sum_{i=1}^M\nabla_\theta\log\pi_{\theta_k}(a_t\mid s_t)\gamma^t(Q^{\pi_\theta}(s_t,a_t)-V^{\phi}(s_t))
    • Update θk\theta_k using g^\hat{g} with an optimizer.

Overview

最后,我们把前面的 basline 使用统一的公式进行表示,记

θJ(θ)=Eτ(p0,πθ,p)[Ψ(τ)θlogpθ(τ)]=Eτ(p0,πθ,p)[t=0TΨt(τ)θlogpθ(τ)]\nabla_\theta \mathcal{J}(\theta)=\mathbb{E}_{\tau\sim(p_0,\pi_\theta, p)}\left[\Psi(\tau)\nabla_\theta \log p_\theta(\tau)\right]=\mathbb{E}_{\tau\sim(p_0,\pi_\theta, p)}\left[\sum_{t=0}^T\Psi_t(\tau)\nabla_\theta \log p_\theta(\tau)\right]

则不同的方法对应的表达式如下

MethodExpression
VallinnaΨt(τ)=G0(τ)=Qπ(st,at)\Psi_t(\tau)=G_0(\tau)=Q^\pi(s_t,a_t)
Return to goΨt(τ)=k=tTγktrk\Psi_t(\tau)=\sum_{k=t}^T\gamma^{k-t}r_k
AdvantageΨt(τ)=Aπ(st,at)=Qπ(st,at)Vπ(st)\Psi_t(\tau)=A^{\pi}(s_t,a_t)=Q^\pi(s_t,a_t)-V^\pi(s_t)

Experiments

TODO

Value-Based Methods v.s. Policy-Based Methods

Value-BasedPolicy-Based
核心思想学习 value/Q-function, 间接得到策略直接优化策略参数
代表算法DQN, Double DQN, Dueling DQNREINFORCE, PPO, TRPO
策略形式隐式 (从 Q 值推导)显式 (参数化策略)
适用场景离散动作空间离散/连续动作空间
收敛性较好需要基线/advantage 降低方差

Actor Critic Methods

上一节使用完整 trajectory return 构造 REINFORCE gradient。这个 estimator 不需要学习 value function,但在长序列和稀疏 reward 下通常具有较高 variance。 baseline invariance 允许我们用 advantage 替代原始 return,而不改变期望梯度。

本节把前面的 policy evaluation 与 policy gradient 合并:actor πθ\pi_\theta 负责生成和更新策略,critic VϕV_\phiQϕQ_\phi 负责估计 return/advantage。对 LLM 而言,两者分别对应需要优化的语言模型和训练期间使用的 value estimator。

Actor-Critic

我们首先考虑最简单的 actor-critic algorithm, 注意到

θJ(θ)=Eτ(p0,πθ,p)[t=0T1θlogπθ(atst)Qπθ(st,at)]\nabla_\theta \mathcal{J}(\theta) = \mathbb{E}_{\tau\sim(p_0,\pi_\theta, p)}\left[ \sum_{t=0}^{T-1}\nabla_\theta\log\pi_\theta(a_t\mid s_t)Q^{\pi_\theta}(s_t,a_t)\right]

因此,其对应的 policy gradient methods 为

θk+1=θk+αEτ(p0,πθ,p)[t=0T1θlogπθ(atst)Qπθ(st,at)]\theta_{k+1} = \theta_k + \alpha \mathbb{E}_{\tau\sim(p_0,\pi_\theta, p)}\left[ \sum_{t=0}^{T-1}\nabla_\theta\log\pi_\theta(a_t\mid s_t)Q^{\pi_\theta}(s_t,a_t)\right]

在前面的章节中,我们介绍了基于 MC 和 TD 两种方式来估计 Qπθ(st,at)Q^{\pi_\theta}(s_t,a_t)。据此可以区分:

  1. 如果 Qπθ(st,at)Q^{\pi_\theta}(s_t,a_t) 由 MC 来进行估计,则我们将其称为 REINFORCE 或者 Monte Carlo policy gradient.
  2. 如果 Qπθ(st,at)Q^{\pi_\theta}(s_t,a_t) 由 TD learning 来进行估计,则我们将其称为 actor-critic, 这是我们本节的重点介绍内容

最简单的 actor-critic algorithm 如下所示

Algorithm: Q Actor-Critic
  1. Initialize policy parameters θ0\theta_0, value function parameters ϕ0\phi_0.
  2. for t=0,1,,T1t = 0, 1, \dots, T-1:
    • atπθ(st)a_t \sim \pi_{\theta}(\cdot \mid s_t), rt,st+1p(,st,at)r_t, s_{t+1} \sim p(\cdot, \cdot \mid s_t, a_t), at+1πθ(st+1)a_{t+1} \sim \pi_{\theta}(\cdot \mid s_{t+1}).
    • (ACTOR) policy update: θt+1=θt+αθθlogπθ(atst)Qϕt(st,at)\theta_{t+1} = \theta_t + \alpha_{\theta} \nabla_{\theta} \log \pi_{\theta}(a_t \mid s_t) \, Q^{\phi_t}(s_t, a_t)
    • (CRITIC) value update: ϕt+1=ϕt+αϕ[rt+γQϕt(st+1,at+1)Qϕt(st,at)]ϕQϕt(st,at)\phi_{t+1} = \phi_t + \alpha_{\phi} \left[r_t + \gamma Q^{\phi_t}(s_{t+1}, a_{t+1}) - Q^{\phi_t}(s_t, a_t)\right] \nabla_{\phi} Q^{\phi_t}(s_t, a_t)

Advantage Actor-Critic (A2C)

在上一节中,我们介绍了 QAC, 但实际上我们用的更多的是 advantage actor-critic, 其梯度如下所示

θJ(θ)=Eτ(p0,πθ,p)[t=0T1θlogπθ(atst)Aπθ(st,at)]\nabla_\theta \mathcal{J}(\theta)=\mathbb{E}_{\tau\sim(p_0,\pi_\theta, p)}\left[ \sum_{t=0}^{T-1}\nabla_\theta\log\pi_\theta(a_t\mid s_t)A^{\pi_\theta}(s_t,a_t)\right]

注意到

Aπθ(st,at)=Qπθ(st,at)Vπθ(st)=Eπθ[rt+γVπθ(st+1)Vπθ(st)st,at]A^{\pi_\theta}(s_t,a_t)=Q^{\pi_\theta}(s_t,a_t)-V^{\pi_\theta}(s_t) = \mathbb{E}^{\pi_\theta}\left[r_t + \gamma V^{\pi_\theta}(s_{t+1})-V^{\pi_\theta}(s_t)\mid s_t, a_t\right]

因此我们可以用 TD error 来进行近似:

A^t=rt+γVϕt(st+1)Vϕt(st)\hat{A}_t= r_t + \gamma V_{\phi_t}(s_{t+1})-V_{\phi_t}(s_t)

这样我们就得到了 A2C 算法

Algorithm: Advantage Actor-Critic (A2C)
  1. Initialize policy parameters θ0\theta_0, value function parameters ϕ0\phi_0.
  2. for t=0,1,,T1t = 0, 1, \dots, T-1:
    • atπθ(st)a_t \sim \pi_{\theta}(\cdot \mid s_t), rt,st+1p(,st,at)r_t, s_{t+1} \sim p(\cdot, \cdot \mid s_t, a_t).
    • Advantage estimation: A^trt+γVϕt(st+1)Vϕt(st)\hat{A}_t \approx r_t + \gamma V_{\phi_t}(s_{t+1}) - V_{\phi_t}(s_t)
    • (ACTOR) policy update: θt+1=θt+αθθlogπθ(atst)A^t\theta_{t+1} = \theta_t + \alpha_{\theta} \nabla_{\theta} \log \pi_{\theta}(a_t \mid s_t) \, \hat{A}_t
    • (CRITIC) value update: ϕt+1=ϕt+αϕA^tϕVϕt(st)\phi_{t+1} = \phi_t + \alpha_{\phi} \, \hat{A}_t \, \nabla_{\phi} V^{\phi_t}(s_t)

GAE

在上一节中,我们介绍了我们可以用 1-step TD error 来近似 advantage,

对比结果如下表所示

Methodbiasvariance
TD errorcan be biasedlow variance
MC estimateunbiasedhigh variance

为了实现更好的 bias-variance trade-off, GAE (Schulman et al., 2016) 提出了使用一个参数 λ[0,1]\lambda\in[0,1] 来控制 bias 和 variance 之间的权重,其表达式如下所示

A^tGAE(γ,λ)==0(γλ)δt+\hat{A}_t^{GAE(\gamma,\lambda)} = \sum_{\ell=0}^{\infty}(\gamma\lambda)^{\ell}\delta_{t+\ell}

其中

δt=rt+γVϕt(st+1)Vϕt(st)\delta_t = r_t + \gamma V_{\phi_t}(s_{t+1})-V_{\phi_t}(s_t)

我们可以探究不同的 λ\lambda 对应的估计:

λ=0\lambda=0 时,我们有

A^tGAE(γ,0)=(γ0)0δt=δt\hat{A}_t^{GAE(\gamma,0)} = (\gamma\cdot0)^0\delta_t=\delta_t

此时,GAE 就退化成了 one-step TD error estimate,

λ=1\lambda=1 时,我们有

A^tGAE(γ,1)==0γδt+==0γ(rt+γVϕt(st+1)Vϕt(st))\hat{A}_t^{GAE(\gamma,1)} = \sum_{\ell=0}^{\infty}\gamma^{\ell}\delta_{t+\ell}=\sum_{\ell=0}^{\infty}\gamma^{\ell}\left(r_t + \gamma V_{\phi_t}(s_{t+1})-V_{\phi_t}(s_t)\right)

如果我们假设这里使用的是真实的 value function 的话,那么我们有

A^tGAE(γ,1)=0γrt++1V(st)=Qπθ(st,at)V(st)\hat{A}_t^{GAE(\gamma,1)} \approx \sum_{\ell=0}^{\infty}\gamma^{\ell}r_{t+\ell+1}-V(s_t)=Q^{\pi_\theta}(s_t,a_t)-V(s_t)

0<λ<10<\lambda<1 时,GAEGAE 是 TD error estimate 与 MC estimate 的一个插值,λ\lambda 越大,estimate 越依赖 long term reward information, 其 bias 越低,但是相应的 variance 越高。λ\lambda 越小,estimate 月依赖于当前的 value function estimate, 其 variance 越低,但是相应的其 bias 越高

实际计算过程中,我们使用如下方法来进行计算

Algorithm: Generalized Advantage Estimation (GAE)
  1. Given policy parameters θ\theta, value function parameters ϕ\phi.
  2. Sample a trajectory τ(p0,πθ,p)\tau \sim (p_0, \pi_{\theta}, p).
  3. for t=0,1,,T1t = 0, 1, \dots, T-1:
    • Compute TD errors: δt=rt+γVϕ(st+1)Vϕ(st)\delta_t = r_t + \gamma V_{\phi}(s_{t+1}) - V_{\phi}(s_t)
  4. Initialize A^TGAE=0\hat{A}_{T}^{\mathrm{GAE}} = 0.
  5. for t=T1,,0t = T-1, \dots, 0:
    • Iterate backward: A^tGAE=δt+γλA^t+1GAE\hat{A}_t^{\mathrm{GAE}} = \delta_t + \gamma\lambda \, \hat{A}_{t+1}^{\mathrm{GAE}}

GAE 的主要优点在于其相对于 MC estimate 可以大幅度降低 variance, 从而提高训练的稳定性以及效率

Why actor critic

comparison with value-based methods:

  1. 无法处理高维或者连续动作空间,actor-critic 中 actor 可以输出一个连续值的概率分布,直接从分布中采样就能得到动作,不需要遍历动作空间
  2. 无法学习到随机策略: value-based methods 只能推导出确定性策略,使用 ϵ\epsilon greedy 之后,也只是随机探索,而不是有策略的随机探索。Actor 可以参数化策略,学习到任意概率分布的随机策略
  3. 学习到的策略更平滑:value-based methods 动作选择依赖于 greedy strategy, 导致 Q function 微小变化导致策略发生剧烈跳变。而 actor-critic 通过 policy gradient 解决了这个问题

comparison with policy-based methods

  1. policy-based methods 依赖完整的 rollout 的 return 来更新,对于 long-term tasks, 最后的 return 方差很大,训练难以收敛。 actor-critic 通过引入 critic 来估计期望收益,计算 advantage, 来降低整体法国差,同时用 TD error 代替回合更新,加速训练。
  1. Schulman, J., Moritz, P., Levine, S., Jordan, M., & Abbeel, P. (2016). High-Dimensional Continuous Control Using Generalized Advantage Estimation. Proceedings of the International Conference on Learning Representations (ICLR).

TRPO

前面的 policy gradient 和 actor-critic 给出了更新方向,但没有回答一个关键问题: 一次可以把 policy 更新多远? 如果新旧 policy 差异过大,旧 policy 生成的 trajectory 很快失效,近似目标也可能不再可靠。

TRPO (Trust Region Policy Optimization) (Schulman et al., 2015) 使用旧策略 πθold\pi_{\theta_{\mathrm{old}}} 的 rollout 构造 surrogate objective,并通过 KL constraint 限制新旧策略的距离。它仍然主要是 on-policy 方法;重点不是把 policy gradient 变成一般的 off-policy learning,而是在一个受控的 trust region 内安全地复用当前批次。

这一思想对 LLM 尤其重要:生成 rollout 成本很高,但在同一批 rollout 上执行过多 gradient steps 又会造成 policy drift。TRPO 首先形式化这一稳定性约束,后续 PPO 会给出更容易实现的一阶近似。

我们的目标函数为

maxπJ(θ)=Es0p0[Vπθ(s0)]\max_{\pi}\mathcal{J}(\theta)=\mathbb{E}_{s_0\sim p_0}\left[V^{\pi_\theta}(s_0)\right]

我们假设有两个 policy 的参数, θ\thetaθold\theta_{old}, 我们分析一下两个 policy 对应的目标函数之间的 difference:

J(θ)J(θold)=J(θ)Es0p0[Vπθold(s0)]=J(θ)Eτ(p0,πθ,p)[Vπθold(s0)]=J(θ)Eτ(p0,πθ,p)[t=0T1γtVπθold(st)t=1T1γtVπθold(st)]=Eτ(p0,πθ,p)[t=0T1γtrt]+Eτ(p0,πθ,p)[t=0T1γt(γVπθold(st+1)Vπθold(st))]=Eτ(p0,πθ,p)[t=0T1γt(rt+γVπθold(st+1)Vπθold(st))]=Eτ(p0,πθ,p)[t=0T1γt(γQπθold(st,at)Vπθold(st))]=Eτ(p0,πθ,p)[t=0T1γtAπθold(st,at)]\begin{aligned} \mathcal{J}(\theta) -\mathcal{J}(\theta_{old}) &= \mathcal{J}(\theta) - \mathbb{E}_{s_0\sim p_0}\left[V^{\pi_{\theta_{old}}}(s_0)\right]\\ &= \mathcal{J}(\theta) - \mathbb{E}_{\tau\sim (p_0,\pi_\theta, p)}\left[V^{\pi_{\theta_{old}}}(s_0)\right]\\ &= \mathcal{J}(\theta) - \mathbb{E}_{\tau\sim (p_0,\pi_\theta, p)}\left[\sum_{t=0}^{T-1}\gamma^tV^{\pi_{\theta_{old}}}(s_t)-\sum_{t=1}^{T-1}\gamma^tV^{\pi_{\theta_{old}}}(s_t)\right]\\ &= \mathbb{E}_{\tau\sim (p_0,\pi_\theta, p)}\left[\sum_{t=0}^{T-1}\gamma^t r_t\right] + \mathbb{E}_{\tau\sim (p_0,\pi_\theta, p)}\left[\sum_{t=0}^{T-1}\gamma^t\left(\gamma V^{\pi_{\theta_{old}}}(s_{t+1})-V^{\pi_{\theta_{old}}}(s_t)\right)\right]\\ &= \mathbb{E}_{\tau\sim (p_0,\pi_\theta, p)}\left[\sum_{t=0}^{T-1}\gamma^t\left(r_t+\gamma V^{\pi_{\theta_{old}}}(s_{t+1})-V^{\pi_{\theta_{old}}}(s_t)\right)\right]\\ &= \mathbb{E}_{\tau\sim (p_0,\pi_\theta, p)}\left[\sum_{t=0}^{T-1}\gamma^t\left(\gamma Q^{\pi_{\theta_{old}}}(s_{t}, a_t)-V^{\pi_{\theta_{old}}}(s_t)\right)\right]\\ &= \mathbb{E}_{\tau\sim (p_0,\pi_\theta, p)}\left[\sum_{t=0}^{T-1}\gamma^t A^{\pi_{\theta_{old}}}(s_t,a_t)\right]\\ \end{aligned}

我们进一步展开得到

J(θ)J(θold)=Eτ(p0,πθ,p)[t=0T1γtAπθold(st,at)]=Es0p0,a0πθ(s0)[Aπθold(s0,a0)+Eπθ[t=1T1γtAπθold(st,at)s0,a0]]=Es0p0,a0πθ(s0),a0πθold(s0)[πθ(a0s0)πθold(a0s0)Aπθold(s0,a0)+Eπθ[t=1T1γtAπθold(st,at)s0,a0]]=(keep enrolling the summation)=Eτ(p0,πθ,p),atπθold(st)[t=0T1γtπθ(atst)πθold(atst)Aπθold(st,at)]\begin{align} \mathcal{J}(\theta) -\mathcal{J}(\theta_{old}) &= \mathbb{E}_{\tau\sim (p_0,\pi_\theta, p)}\left[\sum_{t=0}^{T-1}\gamma^t A^{\pi_{\theta_{old}}}(s_t,a_t)\right]\\ &= \mathbb{E}_{s_0\sim p_0, a_0\sim\pi_\theta(\cdot\mid s_0)}\left[A^{\pi_{\theta_{old}}}(s_0,a_0)+\mathbb{E}^{\pi_\theta}\left[\sum_{t=1}^{T-1}\gamma^t A^{\pi_{\theta_{old}}}(s_t,a_t)\mid s_0, a_0\right]\right]\\ &= \mathbb{E}_{s_0\sim p_0, a_0\sim\pi_\theta(\cdot\mid s_0),a_0'\sim \pi_{\theta_{old}}(\cdot\mid s_0)}\left[\frac{\pi_{\theta}(a_0'\mid s_0)}{\pi_{\theta_{old}}(a_0'\mid s_0)}A^{\pi_{\theta_{old}}}(s_0,a_0')+\mathbb{E}^{\pi_\theta}\left[\sum_{t=1}^{T-1}\gamma^t A^{\pi_{\theta_{old}}}(s_t,a_t)\mid s_0, a_0\right]\right]\\ &=\dots (\text{keep enrolling the summation})\\ &= \mathbb{E}_{\tau\sim (p_0,\pi_\theta, p),a_t'\sim \pi_{\theta_{old}}(\cdot\mid s_t)}\left[\sum_{t=0}^{T-1}\gamma^t \frac{\pi_{\theta}(a_t'\mid s_t)}{\pi_{\theta_{old}}(a_t'\mid s_t)}A^{\pi_{\theta_{old}}}(s_t,a_t')\right]\\ \end{align}

因此,我们有

J(θ)=Eτ(p0,πθ,p),atπθold(st)[t=0T1γtπθ(atst)πθold(atst)Aπθold(st,at)]+C\mathcal{J}(\theta) = \mathbb{E}_{\tau\sim(p_0,\textcolor{red}{\pi_\theta}, p),a_t'\sim\pi_{\theta_{old}}(\cdot\mid s_t)}\left[\sum_{t=0}^{T-1}\gamma^t\frac{\pi_\theta(a_t'\mid s_t)}{\pi_{\theta_{old}}(a_t'\mid s_t)}A^{\pi_{\theta_{old}}}(s_t,a_t')\right]+C

但是,现在我们的更姓还是依赖于 πθ\pi_\theta, 这很难处理。因此 TRPO 的思想就是,πθ\pi_\thetaπθold\pi_{\theta_{old}} 充分接近时,使用 πθold\pi_{\theta_{old}} 来替换 πθ\pi_\theta, 这样,上面目标函数就变成了

maxθ K(θ;θold)=Eτ(p0,πθold,p),atπθold(st)[t=0T1γtπθ(atst)πθold(atst)Aπθold(st,at)]s.t. πθoldπθ\begin{align} \max_{\theta}\ &\mathcal{K}(\theta;\theta_{old}) = \mathbb{E}_{\tau\sim(p_0,\pi_{\theta_{old}}, p),a_t'\sim\pi_{\theta_{old}}(\cdot\mid s_t)}\left[\sum_{t=0}^{T-1}\gamma^t\frac{\pi_\theta(a_t'\mid s_t)}{\pi_{\theta_{old}}(a_t'\mid s_t)}A^{\pi_{\theta_{old}}}(s_t,a_t')\right]\\ \mathrm{s.t.}\ &\pi_{\theta_{old}}\approx \pi_\theta \end{align}

我们可以证明

J(θ)θ=θold=θK(θ;θold)θ=θold\nabla \mathcal{J}(\theta)\mid _{\theta=\theta_{old}} = \nabla_\theta \mathcal{K}(\theta;\theta_{old})\mid _{\theta=\theta_{old}}

这就是 TRPO 的核心改进,在实际求解时,我们将上面的问题规范为如下形式

maxθ K(θ;θold)=Eτ(p0,πθold,p),atπθold(st)[t=0T1γtπθ(atst)πθold(atst)Aπθold(st,at)]s.t. maxsSKL(πθold(s)πθ(s))δ\begin{align} \max_{\theta}\ &\mathcal{K}(\theta;\theta_{old}) = \mathbb{E}_{\tau\sim(p_0,\pi_{\theta_{old}}, p),a_t'\sim\pi_{\theta_{old}}(\cdot\mid s_t)}\left[\sum_{t=0}^{T-1}\gamma^t\frac{\pi_\theta(a_t'\mid s_t)}{\pi_{\theta_{old}}(a_t'\mid s_t)}A^{\pi_{\theta_{old}}}(s_t,a_t')\right]\\ \mathrm{s.t.}\ &\max_{s\in\mathcal{S}}\mathrm{KL}(\pi_{\theta_{old}}(\cdot\mid s)\mid\mid \pi_\theta(\cdot\mid s))\leq \delta \end{align}

这里 δ>0\delta>0 为一个超参数。

现在我们可以采样多条轨迹 τ(i)(p0,πθold,p)\tau^{(i)}\sim (p_0,\pi_{\theta_{old}}, p) 再使用 MC 方法来估计目标函数 K(θ;θold)\mathcal{K}(\theta;\theta_{old})

K(θ;θold)=Eτ(p0,πθold,p),atπθold(st)[t=0T1γtπθ(atst)πθold(atst)Aπθold(st,at)]1Ni=1Nt=0T(i)1γtπθ(at(i)st(i))πθold(at(i)st(i))A^t\mathcal{K}(\theta;\theta_{old}) = \mathbb{E}_{\tau\sim(p_0,\pi_{\theta_{old}}, p),a_t\sim\pi_{\theta_{old}}(\cdot\mid s_t)}\left[\sum_{t=0}^{T-1}\gamma^t\frac{\pi_\theta(a_t\mid s_t)}{\pi_{\theta_{old}}(a_t\mid s_t)}A^{\pi_{\theta_{old}}}(s_t,a_t)\right]\approx \frac{1}{N}\sum_{i=1}^N\sum_{t=0}^{T^{(i)}-1}\gamma^t\frac{\pi_\theta(a_t^{(i)}\mid s_t^{(i)})}{\pi_{\theta_{old}}(a_t^{(i)}\mid s_t^{(i)})}\hat{A}_t

这里 A^tAπθold(st,at)\hat{A}_t\approx A^{\pi_{\theta_{old}}}(s_t,a_t) 是一个 advantage estimator. 注意上面的 estimator 当 NN 变大以及 πθπθold\pi_{\theta}\approx \pi_{\theta_{old}} 时比较准确。

Gamma Trick

在实际实现中,由于 discount factor γ<1\gamma < 1, 远期 advantage 对梯度的贡献会指数级衰减。为了提高效率,TRPO 中常常使用 γ=1\gamma = 1 来简化计算,但此时需要另外一种形式的约束来保证目标函数的有界性。这被称为 gamma trick:通过适当选择 γ\gamma 的值来平衡 bias 和 variance.

更具体地,当 γ<1\gamma < 1 时:

  1. 远期奖励对当前决策的影响被衰减,这在经济学中是合理的(现值概念).
  2. 有助于减少远期估计的方差(远期的 VπV^{\pi} 估计不确定性更大).
  3. 在 TRPO 的 derivation 中,γ\gamma 出现在目标函数的 discount 和中,使用较小的 γ\gamma 可以使 πθπθold\pi_{\theta} \approx \pi_{\theta_{old}} 的近似更准确.

Overview

我们把上面的结果整理为如下算法

Algorithm: TRPO MC
  1. Initialize policy parameters θ0\theta_0, value function parameters ϕ0\phi_0. Set hyperparameter δ>0\delta > 0.
  2. for k=0,1,k = 0, 1, \dots:
    • Sample NN trajectories {τi}\{\tau_i\} according to (p0,πθk,p)(p_0, \pi_{\theta_k}, p).
    • Solve the constrained optimization: maxθk+11Ni=1Nt=0T(i)1πθk+1(at(i)st(i))πθk(at(i)st(i))A^ts.t.maxsSKL(πθk(s)πθk+1(s))δ\begin{aligned} \max_{\theta_{k+1}} \quad &\frac{1}{N} \sum_{i=1}^{N} \sum_{t=0}^{T^{(i)}-1} \frac{\pi_{\theta_{k+1}}(a_t^{(i)} \mid s_t^{(i)})}{\pi_{\theta_k}(a_t^{(i)} \mid s_t^{(i)})} \hat{A}_t \\ \mathrm{s.t.} \quad &\max_{s \in \mathcal{S}} \mathrm{KL}(\pi_{\theta_{k}}(\cdot \mid s) \mid\mid \pi_{\theta_{k+1}}(\cdot \mid s)) \leq \delta \end{aligned}
    • Compute R^t(i)=γttrt(i)\hat{R}_t^{(i)} = \gamma^{t'-t} r_{t'}^{(i)} for i=1,,Ni = 1, \dots, N, t=1,,T(i)1t = 1, \dots, T^{(i)}-1.
    • Update the value model: minϕ1Ni=1N1T(i)t=0T(i)112(Vϕ(st(i))R^t(i))2\min_{\phi} \frac{1}{N} \sum_{i=1}^{N} \frac{1}{T^{(i)}} \sum_{t=0}^{T^{(i)}-1} \frac{1}{2} \left(V_{\phi}(s_t^{(i)}) - \hat{R}_t^{(i)}\right)^2

关于这个优化问题的解法可以参考后续 discussion 章节

Experiments

Discussion

Connection with TRO

我们将目标函数和约束进行泰勒展开得到

K(θ;θold)gT(θθk)KL(θθk)12(θθk)TH(θθk)\begin{align} \mathcal{K}(\theta;\theta_{old}) &\approx g^T(\theta-\theta_k) \\ \mathrm{KL}(\theta\mid\mid \theta_k)&\approx \frac12(\theta-\theta_k)^TH(\theta-\theta_k) \end{align}

这样我们的优化目标就近似为

maxθ gT(θθk)s.t. 12(θθk)TH(θθk)δ\begin{align} \max_{\theta}\ & g^T(\theta-\theta_k)\\ \mathrm{s.t.}\ & \frac12(\theta-\theta_k)^TH(\theta-\theta_k)\leq\delta \end{align}

这个目标函数与 TRO (trust region methods) 一致,因此我们可以使用类似的做法来解决,我们可以得到上面问题的准确表达式

θk+1=θk+2δgTH1gH1g\theta_{k+1} = \theta_k + \sqrt{\frac{2\delta}{g^TH^{-1}g}}H^{-1}g

但是,由于我们使用了 Taylor 展开,实际上我们的更新并不一定满足 KL divergence 的约束,TRPO 对此进行了改进,即在更新时加入了 line search 来使得 θk+1\theta_{k+1} 满足 KL divergence 约束。

θk+1=θk+αj2δgTH1gH1g\theta_{k+1} = \theta_k + \alpha^j\sqrt{\frac{2\delta}{g^TH^{-1}g}}H^{-1}g

其中 α(0,1)\alpha\in(0,1) 是 backtracking line search 的参数, jj 是使得 θk+1\theta_{k+1} 满足 KL divergence 约束的最小正整数 kk.

Drawbacks

虽然 TRPO 的理论形式非常简单,但是实际上实现起来很麻烦。其根本原因在于这个约束比较难以满足,原始论文中使用了 line search 来解决这个问题,但是当状态空间大了之后,速度会显著下降,这也是 PPO 的核心贡献之一。

  1. Schulman, J., Levine, S., Moritz, P., Jordan, M. I., & Abbeel, P. (2015). Trust Region Policy Optimization. Proceedings of the 32nd International Conference on Machine Learning (ICML).

PPO

在上一节中,我们介绍了 TRPO 算法,TRPO 算法是一个二阶算法,TRPO 大幅度提高了 sample efficiency, 但是其问题在于计算/优化太过复杂。因此,PPO 就通过 clip 等技巧保留了 TRPO 思想并降低了计算难度。

在 TRPO 算法中,我们的优化问题形式为:

maxθ K(θ;θold)=Eτ(p0,πθold,p),atπθold(st)[t=0T1πθ(atst)πθold(atst)Aπθold(st,at)]s.t. maxsSKL(πθold(s)πθ(s))δ\begin{align} \max_{\theta}\ &\mathcal{K}(\theta;\theta_{old}) = \mathbb{E}_{\tau\sim(p_0,\pi_{\theta_{old}}, p),a_t'\sim\pi_{\theta_{old}}(\cdot\mid s_t)}\left[\sum_{t=0}^{T-1}\frac{\pi_\theta(a_t'\mid s_t)}{\pi_{\theta_{old}}(a_t'\mid s_t)}A^{\pi_{\theta_{old}}}(s_t,a_t')\right]\\ \mathrm{s.t.}\ &\max_{s\in\mathcal{S}}\mathrm{KL}(\pi_{\theta_{old}}(\cdot\mid s)\mid\mid \pi_{\theta}(\cdot\mid s))\leq \delta \end{align}

TRPO 的核心思想为,当 πθ\pi_{\theta}πθold\pi_{\theta_{old}} 比较接近时,我们的目标函数 K(θ;θold)\mathcal{K}(\theta;\theta_{old}) 与 policy gradient 的目标函数 J(θ)\mathcal{J}(\theta) 梯度误差比较小。

PPO 针对这一点核心观点进行了扩展,我们可以通过另一种形式来约束 πθ\pi_{\theta}πθold\pi_{\theta_{old}} 比较接近,即

rt(θ):=πθ(atst)πθold(atst)1r_t(\theta) := \frac{\pi_\theta(a_t'\mid s_t)}{\pi_{\theta_{old}}(a_t'\mid s_t)} \approx 1

基于这个思想,我们可以丢弃 TRPO 的约束,直接对 rt(θ)r_t(\theta) 进行约束,这个约束可以使用 clip\mathrm{clip} 函数来实现,此时优化的目标函数就变成了

J(θ)=E(q,a)D,otπθold(q)[clip(rt(θ),1ϵ,1+ϵ)A^t]\mathcal{J}(\theta) = \mathbb{E}_{(q,a)\sim\mathcal{D},o_{\leq t}\sim \pi_{\theta_{old}}(\cdot\mid q)}\left[ \mathrm{clip}\left(r_t(\theta), 1-\epsilon, 1+\epsilon\right)\hat{A}_t \right]

这里 ϵ>0\epsilon>0 是一个超参数,类似于 TRPO 中的 δ\delta.

clip(x,,r)={, if xx, if <x<rr, if xr\mathrm{clip}(x,\ell, r)=\begin{cases} \ell, &\text{ if }x\leq \ell\\ x, &\text{ if }\ell<x<r\\ r, &\text{ if }x \geq r \end{cases}

现在我们已经通过 clip\mathrm{clip} 抛弃了 TRPO 复杂的约束了。我们来分析一下目标函数的性质,我们将不同的结果总结为下表

A^t\hat{A}_trtr_tobjectivegradientdescription
>0>0r>1+ϵr>1+\epsilon(1+ϵ)A^t(1+\epsilon)\hat{A}_t0好动作,比 old policy 概率更大,不更新
>0>0[1ϵ,1+ϵ][1-\epsilon,1+\epsilon]rtA^tr_t\hat{A}_tθrtA^t\nabla_\theta r_t\hat{A}_t好动作,与 old policy 概率差不多,更新
>0>0r<1ϵr<1-\epsilon(1ϵ)A^t(1-\epsilon)\hat{A}_t0好动作,比 old policy 概率更小,不更新
<0<0r>1+ϵr>1+\epsilon(1+ϵ)A^t(1+\epsilon)\hat{A}_t0坏动作,比 old policy 概率更大,不更新
<0<0[1ϵ,1+ϵ][1-\epsilon,1+\epsilon]rtA^tr_t\hat{A}_tθrtA^t\nabla_\theta r_t\hat{A}_t坏动作,比 old policy 概率差不多,更新
<0<0r<1ϵr<1-\epsilon(1ϵ)A^t(1-\epsilon)\hat{A}_t0坏动作,比 old policy 概率差更小,不更新

从结果中我们可以看出,只有当 rt1r_t\approx 1 时,我们才会更新我们的模型,对于 A^t>0,rt<1ϵ\hat{A}_t>0, r_t<1-\epsilonA^t<0,rt<1ϵ\hat{A}_t<0, r_t<1-\epsilon 这两种应该更新的情况,我们并没有更新,我们的 sample efficiency 很低。为了解决这个问题,我们在 clip\mathrm{clip} 的基础上,进一步加入一个 min\min 函数来进行控制,这样,我们就得到了 PPO 的目标函数:

JPPO(θ)=E(q,a)D,otπθold(q)[min(rt(θ)A^t,clip(rt(θ),1ϵ,1+ϵ)A^t)]\mathcal{J}_{\mathrm{PPO}}(\theta) = \mathbb{E}_{(q,a)\sim\mathcal{D},o_{\leq t}\sim \pi_{\theta_{old}}(\cdot\mid q)}\left[ \min\left(r_t(\theta)\hat{A}_t,\mathrm{clip}\left(r_t(\theta), 1-\epsilon, 1+\epsilon\right)\hat{A}_t\right) \right]

此时,我们再对目标函数进行分析,就得到

A^t\hat{A}_trtr_tobjectivegradientdescription
>0>0>1+ϵ>1+\epsilon(1+ϵ)A^t(1+\epsilon)\hat{A}_t00好动作,比 old policy 概率更大,不需要更新
>0>01+ϵ\leq1+\epsilonrtA^tr_t\hat{A}_tθrtA^t\nabla_\theta r_t\hat{A}_t好动作,与 old policy 概率差不多或更小,需要更新
<0<01ϵ\geq1-\epsilonrtA^tr_t\hat{A}_tθrtA^t\nabla_\theta r_t\hat{A}_t坏动作,比 old policy 概率更小,不需要概率
<0<0<1ϵ<1-\epsilon(1ϵ)A^t(1-\epsilon)\hat{A}_t00坏动作,与 old policy 概率差不多或更高,概率

可以看到,此时 PPO 的目标函数就解决了 sample efficiency 过低的问题了。

PPO in RLHF

其中

rt(θ)=πθ(otq,o<t)πθold(otq,o<t)r_t(\theta) = \frac{\pi_{\theta}(o_t\mid q, o_{< t})}{\pi_{\theta_{old}}(o_t\mid q, o_{< t})}

(q,a)(q,a) 是从数据集 D\mathcal{D} 采样的 QA pair,ϵ>0\epsilon>0 是一个超参数,A^t\hat{A}_ttt 时刻的优势估计 (advantage estimator). 给定 value function VV 以及 reward function RR, A^t\hat{A}_t 通过计算 GAE 得到:

A^tGAE(γ,λ)=k=0(γλ)kδt+k\hat{A}_t^{\mathrm{GAE}(\gamma, \lambda)}=\sum_{k=0}^{\infty}(\gamma\lambda)^k\delta_{t+k}

其中

δk=Rk+γV(sk+1)V(sk),0γ,λ1\delta_k = R_k + \gamma V(s_{k+1})-V(s_k),\quad 0\leq \gamma,\lambda\leq 1

Implementation

TODO

GRPO

GRPO (Group Relative Policy Optimization) 是 DeepSeek 提出的 RL 算法,相比于 PPO,GRPO 不依赖于 value function(因此也不需要 reward model 来训练 critic),而是通过同一 prompt 下的多组输出来估计 advantage.

Motivation

PPO 需要维护一个 value function VϕV_{\phi} 来估计 advantage AtA_t:

A^tGAE=k=0(γλ)kδt+k,δk=rk+γVϕ(sk+1)Vϕ(sk)\hat{A}_t^{\mathrm{GAE}} = \sum_{k=0}^{\infty} (\gamma\lambda)^k \delta_{t+k}, \quad \delta_k = r_k + \gamma V_{\phi}(s_{k+1}) - V_{\phi}(s_k)

但是在 LLM 场景下,训练 value function 会带来额外的计算和内存开销:

  1. 需要额外训练一个与 policy model 规模相当的 critic model
  2. 需要 reward model 提供的 token-level reward(而很多 verifier 只提供 outcome-level reward)

GRPO 的核心思想:用同一 prompt 下的一组输出来估计 baseline(即群体级别的相对奖励),从而避免训练 value function.

Algorithm

给定 QA pair (q,a)(q, a), 从 πθold\pi_{\theta_{old}} 中采样 GG 个输出 {oi}i=1G\{o_i\}_{i=1}^G, 并获取对应的 reward {Ri}i=1G\{R_i\}_{i=1}^G.

Advantage Estimation

GRPO 使用 group-level 归一化来估计 advantage:

A^i,t=Rimean({Ri}i=1G)std({Ri}i=1G)\hat{A}_{i,t} = \frac{R_i - \mathrm{mean}(\{R_i\}_{i=1}^G)}{\mathrm{std}(\{R_i\}_{i=1}^G)}

这里 advantage A^i,t\hat{A}_{i,t} 与时间步 tt 无关(仅与 group 中的 reward 排名相关),这使得 GRPO 能够直接使用 outcome-level reward.

直观理解:同一 prompt 下的 GG 个 response 中,高于平均奖励的 response 获得正的 advantage, 低于平均奖励的 response 获得负的 advantage. 这与 PPO 中 At=QtVtA_t = Q_t - V_t 的思想一致:VtV_t 是”平均意义上的 expected return”, 起到 baseline 的作用.

GRPO Objective:

GRPO 的训练目标与 PPO 类似,但在分组上进行了归一化:

JGRPO(θ)=E(q,a)D,{oi}i=1Gπθold(q)[1Gi=1G1oit=1oimin(ri,t(θ)A^i,t,  clip(ri,t(θ),1ϵ,1+ϵ)A^i,t)]\mathcal{J}_{\mathrm{GRPO}}(\theta) = \mathbb{E}_{(q, a) \sim \mathcal{D}, \, \{o_i\}_{i=1}^G \sim \pi_{\theta_{old}}(\cdot \mid q)}\left[ \frac{1}{G} \sum_{i=1}^G \frac{1}{|o_i|} \sum_{t=1}^{|o_i|} \min\left(r_{i,t}(\theta) \hat{A}_{i,t}, \; \mathrm{clip}\left(r_{i,t}(\theta), 1 - \epsilon, 1 + \epsilon\right) \hat{A}_{i,t}\right) \right]

其中:

ri,t(θ)=πθ(oi,tq,oi,<t)πθold(oi,tq,oi,<t)r_{i,t}(\theta) = \frac{\pi_{\theta}(o_{i,t} \mid q, o_{i,<t})}{\pi_{\theta_{old}}(o_{i,t} \mid q, o_{i,<t})}

是 importance sampling ratio,衡量新旧策略在 token oi,to_{i,t} 处的概率比.

GRPO 中 A^i,t\hat{A}_{i,t} 对所有 token 使用相同的 advantage 值(基于 outcome reward)。这与 PPO 使用 GAE 得到 per-token advantage 不同.

GRPO v.s. PPO

PPOGRPO
Value function需要训练 VϕV_{\phi}不需要
Advantage estimationGAE (per-token)Group-level 归一化 (per-response)
Reward 要求Token-level reward (需要 reward model)Outcome-level reward (verifier 即可)
计算开销更高 (需训练 critic)更低
采样效率较高需要更多样本 (GG 个 response per prompt)
KL 正则化通常使用 KL penalty通常直接使用 reward shaping

KL Divergence Regularization

在实际的 RLHF/GRPO 训练中,通常会在 reward 中减去一个 KL 惩罚项,以防止模型偏离 reference model 太远:

Ri=Rtask(oi)βKL(πθπref)R_i = R_{\mathrm{task}}(o_i) - \beta \, \mathrm{KL}(\pi_{\theta} \| \pi_{\mathrm{ref}})

其中 β>0\beta > 0 是超参数,πref\pi_{\mathrm{ref}} 是 reference model (通常是 SFT 后的模型).

常用的 KL 估计方式包括:

  1. Kullback-Leibler divergence: KL(πθπref)=Eoπθ[logπθ(oq)πref(oq)]\mathrm{KL}(\pi_{\theta} \| \pi_{\mathrm{ref}}) = \mathbb{E}_{o \sim \pi_{\theta}}\left[\log\frac{\pi_{\theta}(o \mid q)}{\pi_{\mathrm{ref}}(o \mid q)}\right]
  2. Unbiased KL estimator: DeepSeek 使用的低方差估计,形式为: KL(πθπref)=πref(oi,tq,oi,<t)πθ(oi,tq,oi,<t)logπref(oi,tq,oi,<t)πθ(oi,tq,oi,<t)1\mathrm{KL}(\pi_{\theta} \| \pi_{\mathrm{ref}}) = \frac{\pi_{\mathrm{ref}}(o_{i,t} \mid q, o_{i,<t})}{\pi_{\theta}(o_{i,t} \mid q, o_{i,<t})} - \log\frac{\pi_{\mathrm{ref}}(o_{i,t} \mid q, o_{i,<t})}{\pi_{\theta}(o_{i,t} \mid q, o_{i,<t})} - 1

Training Pipeline

GRPO 的典型训练流程如下:

Algorithm: GRPO Training
  • Input: Dataset D\mathcal{D} of prompts, reference model πref\pi_{\mathrm{ref}}, group size GG.
  • Initialize: policy πθ=πref\pi_{\theta} = \pi_{\mathrm{ref}}.

for each training iteration:

  1. Sample a batch of prompts {q}\{q\} from D\mathcal{D}.
  2. For each prompt qq, sample GG responses {oi}\{o_i\} from πθ\pi_{\theta}: oiπθ(q),i=1,,Go_i \sim \pi_{\theta}(\cdot \mid q), \quad i = 1, \dots, G
  3. Compute reward for each response: Ri=Rtask(oi)βKL(πθπref)R_i = R_{\mathrm{task}}(o_i) - \beta \, \mathrm{KL}(\pi_{\theta} \| \pi_{\mathrm{ref}})
  4. Compute group-level advantage: A^i=Rimean({R})std({R})\hat{A}_i = \frac{R_i - \mathrm{mean}(\{R\})}{\mathrm{std}(\{R\})}
  5. Update πθ\pi_{\theta} using GRPO objective: L=1Gi=1G1oit=1oimin ⁣(ri,t(θ)A^i,  clip ⁣(ri,t(θ),1ϵ,1+ϵ)A^i)\mathcal{L} = \frac{1}{G} \sum_{i=1}^{G} \frac{1}{|o_i|} \sum_{t=1}^{|o_i|} \min\!\left(r_{i,t}(\theta) \hat{A}_i, \; \mathrm{clip}\!\left(r_{i,t}(\theta), 1 - \epsilon, 1 + \epsilon\right) \hat{A}_i\right)
  6. (Optional) Update πref\pi_{\mathrm{ref}} periodically.

Why GRPO Works for LLM

  1. 无需求解 value function: 在 LLM 场景下,action space (词表) 和 state space (上下文) 极大,训练 value function 是非常困难且昂贵的。GRPO 通过 group 比较自然避免了这一问题.

  2. Outcome-level reward 友好: 许多 LLM 任务(如数学推理、代码生成)只有最终的 outcome reward (答案正确与否)。GRPO 可以直接利用 outcome reward,无需 reward model 提供 token-level reward.

  3. 简单高效: GRPO 的实现相对简单,无需维护 critic network,计算和内存开销更小.