本文围绕一个核心问题展开:
我们是如何使用强化学习来提高大语言模型的能力上限的?
为了解决这个问题,我们需要回答:
强化学习是如何建模的?
强化学习有哪些算法?哪些强化学习算法适用/不适用于大语言模型?这些算法背后的原理是什么?
现代大语言模型使用的强化学习算法是什么?
我们将按照下面的学习路径组织:
Foundations : 我们将从基本的 定义 (MDP, return) 出发,建立 LLM generation 与 RL 的对应关系,再用 Bellman Equation 将长期回报转化为递推问题。
Evaluation and Improvement : 这一部分,我们先回答“当前 policy 有多好” (policy evaluation ), 再介绍 value-based methods , 该部分会解释为什么 Q-learning/DQN 对理解 RL 很重要,但通常不是优化 LLM policy 的最终选择。
Policy Optimization : 这一部分,我们从 policy gradient methods 逐步过渡到 actor-critic methods , 在这个过程中,我们将逐步优化 policy model 和 value model 的训练与使用。这也是后续 PPO 算法的核心。
Modern RL Algorithms for LLM : 我们将从 TRPO 开始, 逐步介绍 PPO , GRPO 所做的改进,到这一步,我们就接触到了现代 RL 算法的核心部分。
前两部分建立后续推导所需的共同语言;如果读者已经熟悉经典 RL,可以先阅读
LLM as an MDP ,再从 Policy Evaluation
进入算法主线。
这个 tutorial 主要参考了以下两位老师的课程和材料:
本章的目标是建立全文统一使用的语言:我们先用 MDP 描述 agent 与 environment 的交互,再定义 trajectory, return 和优化目标。
最后我们把 auto-regressive generation 写成一个 token-level MDP.
后续章节我们会一直使用这里定义的 LLM state, action 和 reward.
RL 的基本思想是让 agent 通过与环境交互,学习能够最大化期望回报的策略。
其定义为:
强化学习是一个通过构建可以与 environment 进行交互的 agent 来解决控制和决策任务的学习框架,交互的方式为 agent 执行 action, 然后 environment 给予反馈。
RL 的执行过程如下所示
RL 的数学建模依赖于 Markov Decision Process, 下面我们先介绍相关概念。
Definition: Markov Decision Process MDP 是一个可以用于描述强化学习的数学模型,形式化为:
状态 (State) : s t ∈ S s_t \in \mathcal{S} s t ∈ S , S \mathcal{S} S 为状态空间。
动作 (Action) : a t ∈ A a_t \in \mathcal{A} a t ∈ A , A \mathcal{A} A 为动作空间。
奖励 (Reward) : r t ∈ R r_t \in \mathbb{R} r t ∈ R , 环境对 agent 执行动作 a t a_t a t 的反馈。
终止时间 (Terminal Time) : T T T ,且定义 s T = ⟨ term ⟩ s_T = \langle\text{term}\rangle s T = ⟨ term ⟩ 为终止状态。
初始状态分布 (Initial State) : s 0 ∼ p 0 s_0 \sim p_0 s 0 ∼ p 0 .
状态转移概率 (Transition Probability) : ( r t , s t + 1 ) ∼ p ( ⋅ , ⋅ ∣ s t , a t ) (r_t, s_{t+1}) \sim p(\cdot, \cdot \mid s_t, a_t) ( r t , s t + 1 ) ∼ p ( ⋅ , ⋅ ∣ s t , a t ) .
马尔可夫性 (Markov Property) :
p ( s t + 1 ∣ s t , a t , … ) = p ( s t + 1 ∣ s t , a t ) , p ( r t ∣ s t , a t , … ) = p ( r t ∣ s t , a t ) . \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}} p ( s t + 1 ∣ s t , a t , … ) = p ( s t + 1 ∣ s t , a t ) , p ( r t ∣ s t , a t , … ) = p ( r t ∣ s t , a t ) .
为了简化,我们采用以下约定:
当 r t r_t r t 完全由 ( s t , a t ) (s_t,a_t) ( s t , a t ) 决定时,记为 r t = r ( s t , a t ) r_t = r(s_t, a_t) r t = r ( s t , a t ) .
状态转移函数是平稳的 (stationary), 即 p t ( r , s ′ ∣ s , a ) = p ( r , s ′ ∣ s , a ) p_t(r,s'\mid s,a) = p(r,s'\mid s,a) p t ( r , s ′ ∣ s , a ) = p ( r , s ′ ∣ s , a ) .
a t a_t a t 由策略 π \pi π 决定,确定性策略记为 a t = π ( s t ) a_t = \pi(s_t) a t = π ( s t ) , 随机策略记为 a t ∼ π ( ⋅ ∣ s t ) a_t \sim \pi(\cdot \mid s_t) a t ∼ π ( ⋅ ∣ s t ) .
通常我们使用神经网络表示策略,即 π = π θ \pi = \pi_\theta π = π θ , 其中 θ \theta θ 为网络参数。
我们用 S t S_t S t , A t A_t A t 代表 t t t 时刻的状态和动作,这是一个随机变量,用 S t = s t S_t=s_t S t = s t , A t = a t A_t=a_t A t = a t 代表当前状态为 s t s_t s t , 当前动作为 a t a_t a t .
接下来,我们介绍 trajectory, 也就是 RL 算法所需要的数据的定义
Definition: Trajectory Trajectory (rollout) 定义为 agent 与环境完整的一次交互过程:
τ = ( s 0 , a 0 , r 0 , s 1 , a 1 , r 1 , … , s T − 1 , a T − 1 , r T − 1 , s T ) \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) τ = ( s 0 , a 0 , r 0 , s 1 , a 1 , r 1 , … , s T − 1 , a T − 1 , r T − 1 , s T ) 其概率分布为
p ( τ ∣ s 0 , π , p ) = p 0 ( s 0 ) ∏ t = 0 T − 1 π ( a t ∣ s t ) p ( r t , s t + 1 ∣ s t , a t ) 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) p ( τ ∣ s 0 , π , p ) = p 0 ( s 0 ) t = 0 ∏ T − 1 π ( a t ∣ s t ) p ( r t , s t + 1 ∣ s t , a t )
轨迹的概率分布由初始状态分布、策略以及环境模型共同决定,记为 τ ∼ ( p 0 , π , p ) \tau \sim (p_0, \pi, p) τ ∼ ( p 0 , π , p ) .
轨迹的概率分布使用了 Markov property .
本教程仅考虑 finite horizon MDP, 即 T < ∞ T < \infty T < ∞ .
接下来,我们介绍 RL 算法的优化目标 (discounted) return.
Definition: Discounted Return 一条 trajectory 上的 (discounted) return 定义为:
R ( τ ) = ∑ t = 0 T − 1 γ t r t R(\tau) = \sum_{t=0}^{T-1} \gamma^t r_t R ( τ ) = t = 0 ∑ T − 1 γ t r t 其中 γ ∈ ( 0 , 1 ] \gamma \in (0, 1] γ ∈ ( 0 , 1 ] 是 discount factor ,用于将未来的奖励折现到当前时刻。
从时刻 t t t 开始的 return 记为:
G t = ∑ k = t T − 1 γ k − t r k G_t = \sum_{k=t}^{T-1} \gamma^{k-t} r_k G t = k = t ∑ T − 1 γ k − t r k G t G_t G t 存在如下递推关系:
G t = r t + γ G t + 1 G_t = r_t + \gamma G_{t+1} G t = r t + γ G t + 1 易知
R ( τ ) = G 0 . R(\tau) = G_0. R ( τ ) = G 0 .
RL 的最终目标为最大化 expected return, 形式化为:
max π E τ ∼ ( p 0 , π , p ) R ( τ ) = max π E τ ∼ ( p 0 , π , p ) [ ∑ t = 0 T − 1 γ t r t ] \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}} π max E τ ∼ ( p 0 , π , p ) R ( τ ) = π max E τ ∼ ( p 0 , π , p ) [ t = 0 ∑ T − 1 γ t r t ]
这里期望 E π \mathbb{E}^\pi E π 对应的随机变量为 a t ∼ π ( ⋅ ∣ s t ) a_t \sim \pi(\cdot \mid s_t) a t ∼ π ( ⋅ ∣ s t ) , ( r t , s t + 1 ) ∼ p ( ⋅ , ⋅ ∣ s t , a t ) (r_t, s_{t+1}) \sim p(\cdot, \cdot \mid s_t, a_t) ( r t , s t + 1 ) ∼ p ( ⋅ , ⋅ ∣ s t , a t ) .
这个目标函数源于强化学习的核心假设
Reward Hypothesis : 所有的目标都可以被描述为最大化 expected return.
由于实际问题中,初始状态 s 0 s_0 s 0 和状态转移函数 p p p 一般保持不变,因此我们一般省略目标函数中的 s 0 s_0 s 0 和 p p p , 将目标函数简化为
max π E τ ∼ π R ( τ ) \max_{\pi} \quad \mathbb{E}_{\tau\sim\pi} \, R(\tau) π max E τ ∼ π R ( τ )
这一节我们介绍如何用 RL 的定义来建模 auto-regressive LLM.
LLM 定义如下:
y t ∼ L L M ( ⋅ ∣ y < t , x ) y_t \sim \mathrm{LLM}(\cdot \mid y_{<t}, x) y t ∼ LLM ( ⋅ ∣ y < t , x )
其中 x x x 是输入的 prompt, y < t y_{<t} y < t 是已经生成的部分 (generated prefix), y t y_t y t 是下一个要生成的 token.
RL 和 LLM 的对应关系如下 (Shen et al., 2026 ) :
为了让后续公式可以被直接计算和实验验证,我们使用一个 finite-horizon
episodic MDP ToyLLM 作为示例。
它保留 autoregressive generation 的核心结构,同时将 vocabulary 和最大生成长度控制在可以枚举的范围内。
给定一个固定 prompt x x x , 最大 response 长度 T T T 和 target response y ∗ = ( y 0 ∗ , … , y L − 1 ∗ ) y^*=(y_0^*,\ldots,y_{L-1}^*) y ∗ = ( y 0 ∗ , … , y L − 1 ∗ ) , 其中 L ≤ T L\leq T L ≤ T .
对于普通 token 词表 V ‾ \overline{\mathcal V} V , 我们令 V = V ‾ ∪ { ⟨ e o s ⟩ } \mathcal V=\overline{\mathcal V}\cup\{\langle\mathrm{eos}\rangle\} V = V ∪ {⟨ eos ⟩} .
时刻 t t t 的 state 是 prompt 和已经生成的 prefix, action 是下一个 token:
s t = ( x , y < t ) , y < t = ( a 0 , … , a t − 1 ) , a t = y t ∈ V . s_t=(x,y_{<t}),\quad
y_{<t}=(a_0,\ldots,a_{t-1}),\quad
a_t=y_t\in\mathcal V. s t = ( x , y < t ) , y < t = ( a 0 , … , a t − 1 ) , a t = y t ∈ V .
普通 token 会被确定性地追加到 prefix;生成 <eos> 或已经生成 T T T 个普通 token 时,episode 终止。
去掉 <eos> 后的完整 response 记为 y y y ,环境只在 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} R ( y ) = { 1 , 0 , if y = y ∗ , otherwise .
对这个固定 prompt, action space 和非终止状态空间分别为
A = V , S n o n − t e r m i n a l = ⋃ t = 0 T − 1 { x } × V ‾ t . \mathcal A=\mathcal V,\qquad \mathcal S_{\mathrm{non-terminal}}=\bigcup_{t=0}^{T-1}\{x\}\times\overline{\mathcal V}^{\,t}. A = V , S non − terminal = t = 0 ⋃ T − 1 { x } × V t .
令 K = ∣ V ‾ ∣ K=|\overline{\mathcal V}| K = ∣ V ∣ .
如果所有结束后的 response 共用一个 terminal state, 则状态总数为
S : = ∣ S ∣ = 1 + ∑ t = 0 T − 1 K t . S:=|\mathcal S|=1+\sum_{t=0}^{T-1}K^t. S := ∣ S ∣ = 1 + t = 0 ∑ T − 1 K t .
因此,即使 transition 是确定性的,prefix state space 仍然随 horizon 指数增长。
真实 LLM 的 vocabulary 通常包含数万到数十万个 token, 使 action space 也非常大。
ToyLLM 还有一个重要性质:初始状态的 value 存在解析解。令成功所需的 action sequence 为
a ∗ = { ( y 0 ∗ , … , y L − 1 ∗ , ⟨ e o s ⟩ ) , L < T , ( y 0 ∗ , … , y T − 1 ∗ ) , 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} a ∗ = { ( y 0 ∗ , … , y L − 1 ∗ , ⟨ eos ⟩) , ( y 0 ∗ , … , y T − 1 ∗ ) , L < T , L = T ,
并令 m = ∣ a ∗ ∣ m=|a^*| m = ∣ a ∗ ∣ . 如果 reward 在最后一个 action 后产生,则
V π ( s 0 ) = γ m − 1 ∏ t = 0 m − 1 π ( a t ∗ ∣ x , a < t ∗ ) . V^\pi(s_0)
=\gamma^{m-1}\prod_{t=0}^{m-1}
\pi(a_t^*\mid x,a_{<t}^*). V π ( s 0 ) = γ m − 1 t = 0 ∏ m − 1 π ( a t ∗ ∣ x , a < t ∗ ) .
这个解析解可以作为后续 matrix solve, Monte Carlo 和 temporal-difference learning 实现的 ground truth.
若 policy 接近 uniform distribution, 成功概率大致按 ∣ V ∣ − m |\mathcal V|^{-m} ∣ V ∣ − m 下降,也直接展示了 terminal verifier reward 带来的 sparse-reward 问题。
在 Policy Evaluation 一章中,我们将比较依赖完整环境模型的精确解法、只依赖 rollout 的 MC/TD 方法,以及能够在不同 prefix 之间共享信息的 function approximation.
RL 的基础是 MDP, 其核心假设是 Markov property, 即未来的状态和奖励只与当前状态和动作有关,与历史无关。
trajectory return 是随机变量;RL 目标是最大化它在 policy 和环境随机性下的期望。
从 RL 角度看,LLM 本身是一个 policy, 其初始状态为 prompt, 动作空间为 vocabulary 中的 token index, 奖励由多方面给定,environment 仅对 token 进行简单拼接。
ToyLLM 将作为贯穿示例,帮助我们把 Bellman equation、policy evaluation 和 policy
optimization 映射回 token generation。
虽然 action space 和 state space 是离散的,但是其数量级非常大,因此我们通常使用函数逼近来近似 value function 和 Q-function.
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 将 return 的递推结构转化为 value function 的递推结构。
本章我们先研究固定 policy 的 value function, 再研究 optimal policy 中的 optimal value function, 前者通向 policy evaluation, 后者通向 value based methods.
在强化学习中,我们关心的是 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} long-term return = immediate reward + discounted future return
首先,我们介绍两个基本概念,分别是 value function V π ( s t ) V^{\pi}(s_t) V π ( s t ) 和 Q function Q π ( s t , a t ) Q^{\pi}(s_t, a_t) Q π ( s t , a t ) .
Definition: Value Function 我们定义 value function 如下:
V π ( s t ) = E π [ G t ∣ S t = s t ] V^{\pi}(s_t) = \mathbb{E}^{\pi}[G_t \mid S_t = s_t] V π ( s t ) = E π [ G t ∣ S t = s t ] value function 的具体含义为:agent 从当前状态 s t s_t s t 出发,一直遵循当前策略 π \pi π ,最后获取到的 return.
Definition: Q Function Q function (state-action value function) 定义如下:
Q π ( s t , a t ) = E π [ G t ∣ S t = s t , A t = a t ] Q^{\pi}(s_t, a_t) = \mathbb{E}^{\pi}[G_t \mid S_t=s_t, A_t=a_t] Q π ( s t , a t ) = E π [ G t ∣ S t = s t , A t = a t ] 其具体含义为:agent 从当前状态 s t s_t s t 出发,先执行 action a t a_t a t ,再遵循当前策略 π \pi π ,最后获取到的 return.
为了方便,我们令 V π ( ⟨ term ⟩ ) = 0 V^{\pi}(\langle\text{term}\rangle) = 0 V π (⟨ term ⟩) = 0 , Q π ( ⟨ term ⟩ , a ) = 0 , ∀ a ∈ A Q^{\pi}(\langle\text{term}\rangle, a) = 0, \forall a \in \mathcal{A} Q π (⟨ term ⟩ , a ) = 0 , ∀ a ∈ A .
Value function 和 Q function 存在如下关系:
Proposition: Relationship between value function and Q function V π ( s t ) = E a ∼ π ( ⋅ ∣ s t ) [ Q π ( s t , a ) ] V^{\pi}(s_t) = \mathbb{E}_{a \sim \pi(\cdot \mid s_t)}[Q^{\pi}(s_t, a)] V π ( s t ) = E a ∼ π ( ⋅ ∣ s t ) [ Q π ( s t , a )] 即 value function 是 Q function 关于 action a t a_t a t 的期望。
Proof: 由全概率公式,我们有:
V π ( s t ) = E π [ G t ∣ S t = s t ] = ∑ a ∈ A E π [ G t ∣ S t = s t , A t = a ] π ( a ∣ s t ) = ∑ a ∈ A Q π ( s t , a ) π ( a ∣ s t ) = E a ∼ π ( ⋅ ∣ s t ) [ Q π ( s t , 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} V π ( s t ) = E π [ G t ∣ S t = s t ] = a ∈ A ∑ E π [ G t ∣ S t = s t , A t = a ] π ( a ∣ s t ) = a ∈ A ∑ Q π ( s t , a ) π ( a ∣ s t ) = E a ∼ π ( ⋅ ∣ s t ) [ Q π ( s t , a )]
Value function 和 Q function 存在如下递推关系:
Proposition: Iterative property of Value function and Q function V π ( s t ) = E a t ∼ π ( ⋅ ∣ s t ) , ( r t , s t + 1 ) ∼ p ( ⋅ , ⋅ ∣ s t , a t ) [ r t + γ V π ( s t + 1 ) ∣ S t = s t ] Q π ( s t , a t ) = E ( r t , s t + 1 ) ∼ p ( ⋅ , ⋅ ∣ s t , a t ) , a t + 1 ∼ π ( ⋅ ∣ s t + 1 ) [ r t + γ Q π ( s t + 1 , a t + 1 ) ∣ S t = s t , A t = a t ] \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}
} V π ( s t ) Q π ( s t , a t ) = E a t ∼ π ( ⋅ ∣ s t ) , ( r t , s t + 1 ) ∼ p ( ⋅ , ⋅ ∣ s t , a t ) [ r t + γ V π ( s t + 1 ) ∣ S t = s t ] = E ( r t , s t + 1 ) ∼ p ( ⋅ , ⋅ ∣ s t , a t ) , a t + 1 ∼ π ( ⋅ ∣ s t + 1 ) [ r t + γ Q π ( s t + 1 , a t + 1 ) ∣ S t = s t , A t = a t ]
Proof: 我们以 value function 为例:
V π ( s t ) = E π [ G t ∣ S t = s t ] = E π [ r t + γ G t + 1 ∣ S t = s t ] = E a t ∼ π ( ⋅ ∣ s t ) , ( r t , s t + 1 ) ∼ p ( ⋅ , ⋅ ∣ s t , a t ) [ E π [ r t + γ G t + 1 ∣ s t , a t , r t , s t + 1 ] ∣ S t = s t ] = E a t ∼ π ( ⋅ ∣ s t ) , ( r t , s t + 1 ) ∼ p ( ⋅ , ⋅ ∣ s t , a t ) [ r t + γ E π [ G t + 1 ∣ s t + 1 ] ∣ S t = s t ] = E a t ∼ π ( ⋅ ∣ s t ) , ( r t , s t + 1 ) ∼ p ( ⋅ , ⋅ ∣ s t , a t ) [ r t + γ V π ( s t + 1 ) ∣ S t = s t ] \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} V π ( s t ) = E π [ G t ∣ S t = s t ] = E π [ r t + γ G t + 1 ∣ S t = s t ] = E a t ∼ π ( ⋅ ∣ s t ) , ( r t , s t + 1 ) ∼ p ( ⋅ , ⋅ ∣ s t , a t ) [ E π [ r t + γ G t + 1 ∣ s t , a t , r t , s t + 1 ] ∣ S t = s t ] = E a t ∼ π ( ⋅ ∣ s t ) , ( r t , s t + 1 ) ∼ p ( ⋅ , ⋅ ∣ s t , a t ) [ r t + γ E π [ G t + 1 ∣ s t + 1 ] ∣ S t = s t ] = E a t ∼ π ( ⋅ ∣ s t ) , ( r t , s t + 1 ) ∼ p ( ⋅ , ⋅ ∣ s t , a t ) [ r t + γ V π ( s t + 1 ) ∣ S t = s t ] 这里第四个等式使用了 Markov property ,即 G t + 1 G_{t+1} G t + 1 在已知 s t + 1 s_{t+1} s t + 1 的条件下与 s t , a t , r t s_t, a_t, r_t s t , a t , r t 无关。
Q function 的证明类似,略过。
上面的递推关系说明了 value function 和 Q function 满足一个递推关系,接下来, Bellman Equation Theorem 证明了满足这个递归关系的函数一定是 value function/Q function.
Theorem: Bellman Equation Theorem (Value Function) 令 π \pi π 为一个策略,假设 γ ∈ ( 0 , 1 ) \gamma \in (0, 1) γ ∈ ( 0 , 1 ) , ∣ S ∣ < ∞ |\mathcal{S}| < \infty ∣ S ∣ < ∞ 以及 ∣ r ∣ ≤ R < ∞ |r| \leq R < \infty ∣ r ∣ ≤ R < ∞ , 几乎处处成立 (a . s . a.s. a . s . ). 那么 π \pi π 对应的 value function V π : S → R V^{\pi}: \mathcal{S} \to \mathbb{R} V π : S → R 存在,且满足 Bellman equation:
V π ( s t ) = E a t ∼ π ( ⋅ ∣ s t ) , ( r t , s t + 1 ) ∼ p ( ⋅ , ⋅ ∣ s t , a t ) [ r t + γ V π ( s t + 1 ) ∣ S t = s t ] \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 π ( s t ) = E a t ∼ π ( ⋅ ∣ s t ) , ( r t , s t + 1 ) ∼ p ( ⋅ , ⋅ ∣ s t , a t ) [ r t + γ V π ( s t + 1 ) ∣ S t = s t ] 反之,如果存在函数 V : S → R V: \mathcal{S} \to \mathbb{R} V : S → R 满足 Bellman equation, 则 V = V π V = V^{\pi} V = V π .
Proof: 由 iterative property 我们知道 value function 满足 Bellman equation,接下来我们证明唯一性.
假设存在函数 V : S → R V: \mathcal{S} \to \mathbb{R} V : S → R 满足 Bellman equation.
我们定义 Bellman 算子 T π \mathcal{T}^{\pi} T π 为:
( T π V ) ( s t ) : = E a t ∼ π ( ⋅ ∣ s t ) , ( r t , s t + 1 ) ∼ p ( ⋅ , ⋅ ∣ s t , a t ) [ r t + γ V ( s t + 1 ) ∣ S t = s t ] \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}} ( T π V ) ( s t ) := E a t ∼ π ( ⋅ ∣ s t ) , ( r t , s t + 1 ) ∼ p ( ⋅ , ⋅ ∣ s t , a t ) [ r t + γ V ( s t + 1 ) ∣ S t = s t ] 我们证明该算子是一个 contraction mapping, 考虑范数 ∥ V ∥ ∞ = max s ∈ S ∣ V ( s ) ∣ \|V\|_{\infty} = \max_{s \in \mathcal{S}} |V(s)| ∥ V ∥ ∞ = max s ∈ S ∣ V ( s ) ∣ :
∣ ( T π V 1 ) ( s t ) − ( T π V 2 ) ( s t ) ∣ = ∣ E a t ∼ π ( ⋅ ∣ s t ) , ( r t , s t + 1 ) ∼ p ( ⋅ , ⋅ ∣ s t , a t ) [ γ V 1 ( s t + 1 ) − γ V 2 ( s t + 1 ) ] ∣ ≤ γ E a t ∼ π ( ⋅ ∣ s t ) , ( r , s t + 1 ) ∼ p ( ⋅ , ⋅ ∣ s t , a t ) ∣ V 1 ( s t + 1 ) − V 2 ( s t + 1 ) ∣ ≤ γ max s ∈ S ∣ V 1 ( s ) − V 2 ( s ) ∣ = γ ∥ V 1 − V 2 ∥ ∞ \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} ∣ ( T π V 1 ) ( s t ) − ( T π V 2 ) ( s t ) ∣ = E a t ∼ π ( ⋅ ∣ s t ) , ( r t , s t + 1 ) ∼ p ( ⋅ , ⋅ ∣ s t , a t ) [ γ V 1 ( s t + 1 ) − γ V 2 ( s t + 1 ) ] ≤ γ E a t ∼ π ( ⋅ ∣ s t ) , ( r , s t + 1 ) ∼ p ( ⋅ , ⋅ ∣ s t , a t ) ∣ V 1 ( s t + 1 ) − V 2 ( s t + 1 ) ∣ ≤ γ s ∈ S max ∣ V 1 ( s ) − V 2 ( s ) ∣ = γ ∥ V 1 − V 2 ∥ ∞ 上式对任意 s ∈ S s \in \mathcal{S} s ∈ S 成立,因此:
∥ T π V 1 − T π V 2 ∥ ∞ ≤ γ ∥ V 1 − V 2 ∥ ∞ \left\|\mathcal{T}^{\pi} V_1 - \mathcal{T}^{\pi} V_2\right\|_{\infty} \leq \gamma \left\|V_1 - V_2\right\|_{\infty} ∥ T π V 1 − T π V 2 ∥ ∞ ≤ γ ∥ V 1 − V 2 ∥ ∞ 由于 γ < 1 \gamma < 1 γ < 1 , 因此 T π \mathcal{T}^{\pi} T π 是一个 contraction mapping.
根据 不动点定理 ,已知 V π V^{\pi} V π 是一个不动点(满足 Bellman equation),而不动点唯一,因此 V = V π V = V^{\pi} V = V π .
对于 Q function, 我们也有相同的结论:
Theorem: Bellman Equation Theorem (Q Function) 令 π \pi π 为一个策略,假设 γ ∈ ( 0 , 1 ) \gamma \in (0, 1) γ ∈ ( 0 , 1 ) , ∣ S ∣ < ∞ |\mathcal{S}| < \infty ∣ S ∣ < ∞ , ∣ A ∣ < ∞ |\mathcal{A}| < \infty ∣ A ∣ < ∞ 以及 ∣ r ∣ ≤ R < ∞ |r| \leq R < \infty ∣ r ∣ ≤ R < ∞ , 几乎处处成立 (a . s . a.s. a . s . ). 那么 π \pi π 对应的 Q function Q π : S × A → R Q^{\pi}: \mathcal{S} \times \mathcal{A} \to \mathbb{R} Q π : S × A → R 存在,且满足 Bellman equation:
Q π ( s t , a t ) = E ( r t , s t + 1 ) ∼ p ( ⋅ , ⋅ ∣ s t , a t ) , a t + 1 ∼ π ( ⋅ ∣ s t + 1 ) [ r t + γ Q π ( s t + 1 , a t + 1 ) ∣ S t = s t , A t = a t ] \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 t , a t ) = E ( r t , s t + 1 ) ∼ p ( ⋅ , ⋅ ∣ s t , a t ) , a t + 1 ∼ π ( ⋅ ∣ s t + 1 ) [ r t + γ Q π ( s t + 1 , a t + 1 ) ∣ S t = s t , A t = a t ] 反之,如果存在函数 Q : S × A → R Q: \mathcal{S} \times \mathcal{A} \to \mathbb{R} Q : S × A → R 满足 Bellman equation, 则 Q = Q π Q = Q^{\pi} Q = Q π .
Proof: 证明与 value function 的证明类似,略过。
上一节我们介绍了针对一般 policy 的 Bellman equation, 这一节我们给出针对 optimal policy 的 Bellman Optimal Equation.
我们首先给出 optimal policy 的定义
Definition: Optimal Policy 如果策略 π ∗ \pi^* π ∗ 满足:
V π ∗ ( s ) ≥ V π ( s ) , ∀ s ∈ S , ∀ π V^{\pi^*}(s) \geq V^{\pi}(s), \quad \forall s \in \mathcal{S}, \, \forall \pi V π ∗ ( s ) ≥ V π ( s ) , ∀ s ∈ S , ∀ π 则称策略 π ∗ \pi^* π ∗ 是 optimal policy .
对应的 V π ∗ V^{\pi^*} V π ∗ 和 Q π ∗ Q^{\pi^*} Q π ∗ 分别称为 optimal value function 和 optimal Q function , 简记为 V ∗ = V π ∗ V^* = V^{\pi^*} V ∗ = V π ∗ , Q ∗ = Q π ∗ Q^* = Q^{\pi^*} Q ∗ = Q π ∗ .
⚠️
optimal policy 通常依赖于状态 s s s ,并且可能不唯一;但所有 optimal policy 对应同一个 optimal value function V ∗ V^* V ∗ .
与 value function V π V^\pi V π 和 Q function Q π Q^\pi Q π 之间存在联系一样,optimal value function V ∗ V^{*} V ∗ 和 optimal Q function Q ∗ Q^{*} Q ∗ 也存在类似关系
Proposition: Relationship between optimal value function and optimal Q function V ∗ ( s t ) = max a ∈ A Q ∗ ( s t , a ) \boxed{V^*(s_t) = \max_{a \in \mathcal{A}} Q^*(s_t, a)} V ∗ ( s t ) = a ∈ A max Q ∗ ( s t , a ) 即 V ∗ V^* V ∗ 就是 Q ∗ Q^* Q ∗ 的最优动作对应的 return.
Proof: 我们先证左边小于等于右边:
V ∗ ( s t ) = E a ∼ π ∗ ( ⋅ ∣ s t ) [ Q ∗ ( s t , a ) ] ≤ E a ∼ π ∗ ( ⋅ ∣ s t ) [ max a ∈ A Q ∗ ( s t , a ) ] = max a ∈ A Q ∗ ( s t , 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) V ∗ ( s t ) = E a ∼ π ∗ ( ⋅ ∣ s t ) [ Q ∗ ( s t , a )] ≤ E a ∼ π ∗ ( ⋅ ∣ s t ) [ a ∈ A max Q ∗ ( s t , a )] = a ∈ A max Q ∗ ( s t , a ) 再证右边小于等于左边:
令 a ∗ ∈ arg max a Q ∗ ( s t , a ) a^* \in \arg\max_{a} Q^*(s_t, a) a ∗ ∈ arg max a Q ∗ ( s t , a ) , 令策略 π ′ \pi' π ′ 为确定性策略 π ′ ( a ∗ ∣ s t ) = 1 \pi'(a^* \mid s_t) = 1 π ′ ( a ∗ ∣ s t ) = 1 , 则:
max a ∈ A Q ∗ ( s t , a ∗ ) = Q π ′ ( s t , a ∗ ) = V π ′ ( s t ) ≤ V ∗ ( s t ) \max_{a \in \mathcal{A}} Q^*(s_t, a^*) = Q^{\pi'}(s_t, a^*) = V^{\pi'}(s_t) \leq V^*(s_t) a ∈ A max Q ∗ ( s t , a ∗ ) = Q π ′ ( s t , a ∗ ) = V π ′ ( s t ) ≤ V ∗ ( s t ) 因此两边相等。
接下来,我们给出关于 optimal value function V ∗ V^* V ∗ 对应的 Bellman optimality equation.
Theorem: Bellman Optimality Equation Theorem (value function) 假设 γ ∈ ( 0 , 1 ) \gamma \in (0, 1) γ ∈ ( 0 , 1 ) , ∣ S ∣ < ∞ |\mathcal{S}| < \infty ∣ S ∣ < ∞ , ∣ A ∣ < ∞ |\mathcal{A}| < \infty ∣ A ∣ < ∞ 以及 ∣ r ∣ ≤ R < ∞ |r| \leq R < \infty ∣ r ∣ ≤ R < ∞ , 几乎处处成立 (a.s.). 那么 optimal value function V ∗ : S → R V^*:\mathcal{S} \to \mathbb{R} V ∗ : S → R 存在,且满足 Bellman optimality equation:
V ∗ ( s t ) = max a ∈ A E ( r t , s t + 1 ) ∼ p ( ⋅ , ⋅ ∣ s t , a t ) [ r t + γ V ∗ ( s t + 1 ) ∣ S t = s t , A t = a t ] \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 ∗ ( s t ) = a ∈ A max E ( r t , s t + 1 ) ∼ p ( ⋅ , ⋅ ∣ s t , a t ) [ r t + γ V ∗ ( s t + 1 ) ∣ S t = s t , A t = a t ] 反之,如果存在函数 V : S → R V: \mathcal{S} \to \mathbb{R} V : S → R 满足 Bellman optimality equation, 则 V = V ∗ V = V^* V = V ∗ .
最后,我们可以给出一个 optimal deterministic policy:
π ∗ ( s t ) = arg max a ∈ A E ( r t , s t + 1 ) ∼ p ( ⋅ , ⋅ ∣ s t , a t ) [ r t + γ V ∗ ( s t + 1 ) ∣ S t = s t , A t = a t ] = arg max a ∈ A Q ∗ ( s t , 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}} π ∗ ( s t ) = arg a ∈ A max E ( r t , s t + 1 ) ∼ p ( ⋅ , ⋅ ∣ s t , a t ) [ r t + γ V ∗ ( s t + 1 ) ∣ S t = s t , A t = a t ] = arg a ∈ A max Q ∗ ( s t , a )
证明需要用到下面的引理
Lemma: Monotonicity of Bellman operator 令 π \pi π 为一个策略,T π \mathcal{T}^{\pi} T π 和 T ∗ \mathcal{T}^* T ∗ 分别是 Bellman operator 和 Bellman optimality operator, 其中 Bellman optimality operator 定义为
( T ∗ V ) ( s t ) : = max a ∈ A E ( r t , s t + 1 ) ∼ p ( ⋅ , ⋅ ∣ s t , a t ) [ r t + γ V ( s t + 1 ) ∣ S t = s t , A t = a t ] \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}} ( T ∗ V ) ( s t ) := a ∈ A max E ( r t , s t + 1 ) ∼ p ( ⋅ , ⋅ ∣ s t , a t ) [ r t + γ V ( s t + 1 ) ∣ S t = s t , A t = a t ] 我们有:
对任意 V : S → R V: \mathcal{S} \to \mathbb{R} V : S → R , 有 T π V ( s ) ≤ T ∗ V ( s ) , ∀ s ∈ S \mathcal{T}^{\pi} V(s) \leq \mathcal{T}^* V(s), \, \forall s \in \mathcal{S} T π V ( s ) ≤ T ∗ V ( s ) , ∀ s ∈ S .
对任意 U , V : S → R U, V: \mathcal{S} \to \mathbb{R} U , V : S → R , 如果 U ≤ V U \leq V U ≤ V , 则 T ∗ U ( s ) ≤ T ∗ V ( s ) , ∀ s ∈ S \mathcal{T}^* U(s) \leq \mathcal{T}^* V(s), \, \forall s \in \mathcal{S} T ∗ U ( s ) ≤ T ∗ V ( s ) , ∀ s ∈ S .
证明如下
Proof: 我们首先证明 Bellman optimality operator T ∗ \mathcal{T}^* T ∗ 是一个 contraction mapping: 对任意 V 1 , V 2 V_1, V_2 V 1 , V 2 :
∣ ( T ∗ V 1 ) ( s t ) − ( T ∗ V 2 ) ( s t ) ∣ = ∣ max a E [ r t + γ V 1 ( s t + 1 ) ] − max a E [ r t + γ V 2 ( s t + 1 ) ] ∣ ≤ max a ∣ E [ r t + γ V 1 ( s t + 1 ) ] − E [ r t + γ V 2 ( s t + 1 ) ] ∣ = γ max a ∣ E [ V 1 ( s t + 1 ) − V 2 ( s t + 1 ) ] ∣ ≤ γ ∥ V 1 − V 2 ∥ ∞ \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} ∣ ( T ∗ V 1 ) ( s t ) − ( T ∗ V 2 ) ( s t ) ∣ = a max E [ r t + γ V 1 ( s t + 1 )] − a max E [ r t + γ V 2 ( s t + 1 )] ≤ a max ∣ E [ r t + γ V 1 ( s t + 1 )] − E [ r t + γ V 2 ( s t + 1 )] ∣ = γ a max ∣ E [ V 1 ( s t + 1 ) − V 2 ( s t + 1 )] ∣ ≤ γ ∥ V 1 − V 2 ∥ ∞ 这里第一个不等式使用了 ∣ max s v ( s ) − max s u ( s ) ∣ ≤ max s ∣ u ( s ) − v ( s ) ∣ |\max_s v(s) - \max_s u(s)| \leq \max_s |u(s) - v(s)| ∣ max s v ( s ) − max s u ( s ) ∣ ≤ max s ∣ u ( s ) − v ( s ) ∣ .
因此 ∥ T ∗ V 1 − T ∗ V 2 ∥ ∞ ≤ γ ∥ V 1 − V 2 ∥ ∞ \|\mathcal{T}^* V_1 - \mathcal{T}^* V_2\|_{\infty} \leq \gamma \|V_1 - V_2\|_{\infty} ∥ T ∗ V 1 − T ∗ V 2 ∥ ∞ ≤ γ ∥ V 1 − V 2 ∥ ∞ , T ∗ \mathcal{T}^* T ∗ 是 contraction mapping.
根据不动点定理 ,T ∗ \mathcal{T}^* T ∗ 存在唯一不动点 V ∗ V^* V ∗ .
我们定义:
π ∗ ( s t ) = arg max a ∈ A E ( r t , s t + 1 ) ∼ p ( ⋅ , ⋅ ∣ s t , a t ) [ r + γ V ∗ ( s t + 1 ) ∣ S t = s t , A t = a t ] \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] π ∗ ( s t ) = arg a ∈ A max E ( r t , s t + 1 ) ∼ p ( ⋅ , ⋅ ∣ s t , a t ) [ r + γ V ∗ ( s t + 1 ) ∣ S t = s t , A t = a t ] 此时,我们有 TODO
V ∗ ( s t ) = E π ∗ [ r t + γ V ∗ ( s t + 1 ) ∣ s t ] = E π ∗ [ r t + γ E π ∗ [ r t + 1 + γ V ∗ ( s t + 2 ) ∣ s t + 1 ] ∣ s t ] = E π ∗ [ E π ∗ [ r t + γ ( r t + 1 + γ V ∗ ( s t + 2 ) ) ∣ s t + 1 ] ∣ s t ] = E π ∗ [ r t + γ r t + 1 + γ 2 V ∗ ( s t + 2 ) ∣ s t ] = E π ∗ [ r t + γ r t + 1 + γ 2 r t + 2 + ⋯ ∣ s t ] = V π ∗ ( s t ) \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 ∗ ( s t ) = E π ∗ [ r t + γ V ∗ ( s t + 1 ) ∣ s t ] = E π ∗ [ r t + γ E π ∗ [ r t + 1 + γ V ∗ ( s t + 2 ) ∣ s t + 1 ] ∣ s t ] = E π ∗ [ E π ∗ [ r t + γ ( r t + 1 + γ V ∗ ( s t + 2 )) ∣ s t + 1 ] ∣ s t ] = E π ∗ [ r t + γ r t + 1 + γ 2 V ∗ ( s t + 2 ) ∣ s t ] = E π ∗ [ r t + γ r t + 1 + γ 2 r t + 2 + ⋯ ∣ s t ] = V π ∗ ( s t ) 即 V ∗ = V π ∗ V^*=V^{\pi^*} V ∗ = V π ∗ , 这里第四个等式使用了 Law of Total Expectation .
现在我们证明 π ∗ \pi^* π ∗ 是最优策略,令 π \pi π 为任意一个策略,我们有
V π = T π V ≤ T ∗ V ≤ ( T ∗ ) 2 ( V π ) ≤ ⋯ ≤ ( T ∗ ) k ( V π ) → k → ∞ fixed-point / contraction V ∗ . 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 π = T π V ≤ T ∗ V ≤ ( T ∗ ) 2 ( V π ) ≤ ⋯ ≤ ( T ∗ ) k ( V π ) fixed-point / contraction k → ∞ V ∗ . 即 V π ≤ V ∗ V^\pi\leq V^* V π ≤ V ∗ , 这里第一个等式和第二个等式利用了 Lemma , 由于 V π V^\pi V π 是任意的,我们有 V ∗ = V π ∗ V^*=V^{\pi^*} V ∗ = V π ∗ .
对于 optimal Q function Q ∗ Q^* Q ∗ , 我们也有类似的结论。
Theorem: Bellman Optimality Equation Theorem (Q function) 假设 γ ∈ ( 0 , 1 ) \gamma \in (0, 1) γ ∈ ( 0 , 1 ) , ∣ S ∣ < ∞ |\mathcal{S}| < \infty ∣ S ∣ < ∞ , ∣ A ∣ < ∞ |\mathcal{A}| < \infty ∣ A ∣ < ∞ 以及 ∣ r ∣ ≤ R < ∞ |r| \leq R < \infty ∣ r ∣ ≤ R < ∞ , 几乎处处成立 (a.s.). 那么 optimal Q function Q ∗ : S × A → R Q^*: \mathcal{S} \times \mathcal{A} \to \mathbb{R} Q ∗ : S × A → R 存在,且满足 Bellman optimality equation:
Q ∗ ( s t , a t ) = E ( r t , s t + 1 ) ∼ p ( ⋅ , ⋅ ∣ s t , a t ) [ r + γ max a t + 1 ∈ A Q ∗ ( s t + 1 , a t + 1 ) ∣ S t = s t , A t = a t ] \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 t , a t ) = E ( r t , s t + 1 ) ∼ p ( ⋅ , ⋅ ∣ s t , a t ) [ r + γ a t + 1 ∈ A max Q ∗ ( s t + 1 , a t + 1 ) ∣ S t = s t , A t = a t ] 反之,如果存在函数 Q : S × A → R Q: \mathcal{S} \times \mathcal{A} \to \mathbb{R} Q : S × A → R 满足 Bellman optimality equation, 则 Q = Q ∗ Q = Q^* Q = Q ∗ .
Proof: 证明与 value function 类似,略过。
下一章我们将从这个精确但不可扩展的起点出发,逐步引入 sampling 和 function approximation 来完成 policy evaluation.
G t = r t + γ G t + 1 G_t=r_t+\gamma G_{t+1} G t = r t + γ G t + 1 使 return 可以递归分解;对它取条件期望便得到 Bellman equation.
Bellman equation 回答“固定 policy 的价值是什么”,Bellman optimality equation 回答“所有 policy 中的最优价值是什么”。
对 LLM,value 描述从 (prompt, token prefix) 继续生成后的期望最终得分; Bellman recursion 则承担将稀疏 terminal reward 传回早期 token 的 credit assignment.
在 Bellman Equation Theorem 一节中,我们看到 V π V^\pi V π 是 Bellman operator T π \mathcal T^\pi T π 的唯一不动点。
Bellman optimality equation 直接寻找最优价值,而本章先解决一个更基础的 evaluation 问题:
给定一个策略 π \pi π ,如何高效求解策略 π \pi π 对应的 value function V π V^\pi V π ?
我们将先从完整 MDP 模型的精确解法,逐步过渡到只依赖 rollout 的 MC/TD 方法和现代 RL 使用的函数逼近方法。
这条路径同时回答两个问题:
为什么 Bellman equation 在小环境中可直接求解
精确求解方法对于 LLM 来说存在什么问题
这一节我们介绍针对离散小空间场景下的求解和优化方案。
对于离散状态空间来说,我们可以将 value iteration 算法中的 Bellman Equation 写为如下的矩阵形式
V π = R π + γ P π V π V_{\pi} = R_{\pi} + \gamma P_{\pi} V_{\pi} V π = R π + γ P π V π
其中
V π = [ V π ( s 1 ) , … , V π ( s S ) ] T ∈ R S R π = [ R π ( s 1 ) , … , R π ( s S ) ] T ∈ R S P π = [ P π ( s j ∣ s i ) ] ∈ R S × 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} V π R π P π = [ V π ( s 1 ) , … , V π ( s S ) ] T ∈ R S = [ R π ( s 1 ) , … , R π ( s S ) ] T ∈ R S = [ P π ( s j ∣ s i )] ∈ R S × S
可以看到,上面这个其实是一个线性方程组,即
( I S − γ P π ) V π = R π (\mathbf{I}_S - \gamma P_{\pi})V_{\pi} = R_{\pi} ( I S − γ P π ) V π = R π
对于这个线性方程组,我们可以用矩阵求逆或者迭代解法求解。
精确算法如下所示
Algorithm: Value function Matrix Exact Solve Input : policy π \pi π , discount factor γ \gamma γ , state space S \mathcal{S} S with size S S S
assign index i ( s ) i(s) i ( s ) for s ∈ S , i = 1 , … , S s\in\mathcal{S}, i=1,\dots,S s ∈ S , i = 1 , … , S
P = 0 S × S P=\mathbf{0}_{S\times S} P = 0 S × S , R = 0 S R=\mathbf{0}_{S} R = 0 S
For each s ∈ S s\in\mathcal{S} s ∈ S
i ← i ( s ) i\gets i(s) i ← i ( s )
For each a ∈ A a\in\mathcal{A} a ∈ A
( s ′ , r , d o n e ) ← S T E P ( s , a ) (s', r, \mathrm{done}) \gets \mathrm{STEP}(s, a) ( s ′ , r , done ) ← STEP ( s , a )
R [ i ] ← R [ i ] + π ( a ∣ s ) ⋅ r R[i]\gets R[i]+\pi(a\mid s)\cdot r R [ i ] ← R [ i ] + π ( a ∣ s ) ⋅ r
If NOT d o n e \mathrm{done} done
j ← i ( s ′ ) j\gets i(s') j ← i ( s ′ )
P [ i , j ] ← P [ i , j ] + π ( a ∣ s ) P[i,j]\gets P[i,j]+\pi(a\mid s) P [ i , j ] ← P [ i , j ] + π ( a ∣ s )
V ← S O L V E ( I S − γ P , R ) V\gets \mathrm{SOLVE}(I_S-\gamma P, R) V ← SOLVE ( I S − γ P , R )
return V V V
迭代算法则是给定随机初始化的 V 0 V_0 V 0 , 我们通过迭代计算得到序列 { V k } \{V_k\} { V k } , 其中 V k + 1 = R π + γ P π V k V_{k+1}=R_{\pi} + \gamma P_{\pi}V_k V k + 1 = R π + γ P π V k , 我们可以证明这个序列收敛到 V π V_{\pi} V π .
Algorithm: Value function Matrix Iterative Solve Input : policy π \pi π , discount factor γ \gamma γ , convergence threshold θ \theta θ , state space S \mathcal{S} S
V ( s ) = 0 , s ∈ S V(s)=0, s\in\mathcal{S} V ( s ) = 0 , s ∈ S
While Δ > θ \Delta > \theta Δ > θ :
Δ ← 0 \Delta\gets0 Δ ← 0
For each s ∈ S s\in\mathcal{S} s ∈ S
v ← 0 v\gets0 v ← 0
For each a ∈ A a\in\mathcal{A} a ∈ A
( s ′ , r , d o n e ) ← S T E P ( s , a ) (s', r, \mathrm{done}) \gets \mathrm{STEP}(s, a) ( s ′ , r , done ) ← STEP ( s , a )
If d o n e \mathrm{done} done
v ← v + π ( a ∣ s ) ⋅ r v \gets v + \pi(a|s) \cdot r v ← v + π ( a ∣ s ) ⋅ r
Else
v ← v + π ( a ∣ s ) ⋅ ( r + γ ⋅ V ( s ′ ) ) v \gets v + \pi(a|s) \cdot (r + \gamma \cdot V(s')) v ← v + π ( a ∣ s ) ⋅ ( r + γ ⋅ V ( s ′ ))
Δ ← m a x ( Δ , ∣ v − V ( s ) ∣ ) \Delta ← max(\Delta, |v − V(s)|) Δ ← ma x ( Δ , ∣ v − V ( s ) ∣ )
V ( s ) ← v V(s) \gets v V ( s ) ← v
return V V V
下面是两种算法在不同规模下对应的求解时间变化曲线
可以看到,由于稠密线性方程组的直接求解通常需要 O ( S 3 ) \mathcal{O}(S^3) O ( S 3 ) 时间和 O ( S 2 ) \mathcal{O}(S^2) O ( S 2 ) 空间,随着状态数 S S S 的增加,其所需时间急剧上升。
而迭代解法虽然比矩阵直接求解效率更高,但是仍然需要枚举全部状态,当状态空间增大之后,迭代解法的时间也无法接受。
因此,我们需要计算更加高效的算法。
由 V π V^{\pi} V π 定义,我们有:
V π ( s 0 ) = E π [ R ( τ ) ∣ s 0 ] V^{\pi}(s_0) = \mathbb{E}^{\pi}\left[R(\tau) \mid s_0\right] V π ( s 0 ) = E π [ R ( τ ) ∣ s 0 ]
由于 return 是一个随机变量,我们可以利用 Monte Carlo (MC) 方法得到一个 Unbiased Estimator ,即从 s 0 s_0 s 0 出发,独立随机采样 (i.i.d.) M M M 条轨迹:
{ τ ( i ) = ( s 0 , a 0 ( i ) , r 0 ( i ) , … , a T ( i ) − 1 ( i ) , r T ( i ) − 1 ( i ) , s T ( 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 { τ ( i ) = ( s 0 , a 0 ( i ) , r 0 ( i ) , … , a T ( i ) − 1 ( i ) , r T ( i ) − 1 ( i ) , s T ( i ) ( i ) )} , i = 1 , … , M
然后我们使用样本平均 (sample mean) 来近似期望:
V π ( s 0 ) = E π [ R ( τ ) ∣ s 0 ] ≈ 1 M ∑ i = 1 M R ( τ ( i ) ) V^{\pi}(s_0) = \mathbb{E}^{\pi}[R(\tau) \mid s_0] \approx \frac{1}{M} \sum_{i=1}^{M} R(\tau^{(i)}) V π ( s 0 ) = E π [ R ( τ ) ∣ s 0 ] ≈ M 1 i = 1 ∑ M R ( τ ( i ) )
对应的算法如下所示
Algorithm: Tabular Learning with MC Input : policy π \pi π , discount factor γ \gamma γ , episodes M M M , state space S \mathcal{S} S , initial state distribution p 0 p_0 p 0
V ( s ) = 0 , R e t u r n S u m ( s ) = 0 , R e t u r n C o u n t ( s ) = 0 , s ∈ S V(s)=0, \mathrm{ReturnSum}(s)=0, \mathrm{ReturnCount}(s)=0, s\in\mathcal{S} V ( s ) = 0 , ReturnSum ( s ) = 0 , ReturnCount ( s ) = 0 , s ∈ S
For e p i s o d e = 1 , 2 , … , M \mathrm{episode}=1,2,\dots,M episode = 1 , 2 , … , M :
s 0 ∼ p 0 s_0\sim p_0 s 0 ∼ p 0 , τ = [ ] \tau=[] τ = [ ]
While not d o n e \mathrm{done} done :
a ∼ π ( ⋅ ∣ s ) a\sim\pi(\cdot\mid s) a ∼ π ( ⋅ ∣ s )
( s ′ , r , d o n e ) ← S T E P ( s , a ) (s', r, \mathrm{done}) \gets \mathrm{STEP}(s, a) ( s ′ , r , done ) ← STEP ( s , a )
τ . a p p e n d ( ( s , r ) ) \tau.\mathrm{append}((s,r)) τ . append (( s , r ))
s ← s ′ s\gets s' s ← s ′
R e t u r n s = [ ] \mathrm{Returns}=[] Returns = [ ] , G ← 0 G\gets0 G ← 0
For ( s , r ) (s,r) ( s , r ) in r e v e r s e d ( τ ) \mathrm{reversed}(\tau) reversed ( τ )
G ← r + γ G G\gets r + \gamma G G ← r + γ G
R e t u r n s . p r e p e n d ( ( s , G ) ) \mathrm{Returns}.\mathrm{prepend}((s,G)) Returns . prepend (( s , G ))
V i s i t e d = ∅ \mathrm{Visited}=\emptyset Visited = ∅
For ( s , G ) (s,G) ( s , G ) in R e t u r n s \mathrm{Returns} Returns
If s ∉ V i s i t e d s\not\in \mathrm{Visited} s ∈ Visited
V i s i t e d . a d d ( s ) \mathrm{Visited}.\mathrm{add}(s) Visited . add ( s )
R e t u r n S u m ( s ) ← R e t u r n S u m ( s ) + G \mathrm{ReturnSum}(s) \gets \mathrm{ReturnSum}(s) +G ReturnSum ( s ) ← ReturnSum ( s ) + G
R e t u r n C o u n t ( s ) ← R e t u r n C o u n t ( s ) + 1 \mathrm{ReturnCount}(s) \gets \mathrm{ReturnCount}(s)+ 1 ReturnCount ( s ) ← ReturnCount ( s ) + 1
V ( s ) ← R e t u r n S u m ( s ) / R e t u r n C o u n t ( s ) V(s)\gets \mathrm{ReturnSum}(s) / \mathrm{ReturnCount}(s) V ( s ) ← ReturnSum ( s ) / ReturnCount ( s )
return V V V
MC 的优点是 estimator 是一个 unbiased estimator, 其缺点在于我们必须完整采样完一条轨迹才能得到这个估计,对于 long-horizon task, MC 方法效率很低。
为了解决这个问题,我们可以使用 Temporal difference (TD) learning 方法。
Temporal difference (TD) learning 通过相邻两步的转换关系来近似 value function.
其核心思想在于
注意到:
V π ( s t ) = E π [ r t + γ V π ( s t + 1 ) ∣ S t = s t ] V^{\pi}(s_t) = \mathbb{E}^{\pi}[r_t + \gamma V^{\pi}(s_{t+1}) \mid S_t=s_t] V π ( s t ) = E π [ r t + γ V π ( s t + 1 ) ∣ S t = s t ]
我们首先定义 TD target 和 TD error 如下
Definition: TD target and TD error 给定当前 estimator V t V_t V t ,我们定义 one-step TD target 为
V ‾ t = { r t , if s t + 1 is terminal , r t + γ V t ( s t + 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} V t = { r t , r t + γ V t ( s t + 1 ) , if s t + 1 is terminal , otherwise . 定义 TD error 为
δ t = V ‾ t − V t ( s t ) . \delta_t = \overline{V}_t - V_t(s_t). δ t = V t − V t ( s t ) .
接下来,我们基于 TD error 更新 V t V_t V t :
V t + 1 ( s ) = { V t ( s t ) + α t ( s t ) δ t , if s = s t , V t ( 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} V t + 1 ( s ) = { V t ( s t ) + α t ( s t ) δ t , V t ( s ) , if s = s t , otherwise .
其中 α t ( s t ) ∈ ( 0 , 1 ) \alpha_t(s_t) \in (0, 1) α t ( s t ) ∈ ( 0 , 1 ) 是 learning rate.
TD learning 对应的算法如下所示
Algorithm: Tabular Learning with TD Input : policy π \pi π , discount factor γ \gamma γ , learning rate α \alpha α , episodes M M M , state space S \mathcal{S} S , initial state distribution p 0 p_0 p 0
V ( s ) = 0 , s ∈ S V(s)=0, s\in\mathcal{S} V ( s ) = 0 , s ∈ S
For e p i s o d e = 1 , 2 , … , M \mathrm{episode}=1,2,\dots,M episode = 1 , 2 , … , M :
s 0 ∼ p 0 s_0\sim p_0 s 0 ∼ p 0
While not d o n e \mathrm{done} done :
a ∼ π ( ⋅ ∣ s ) a\sim\pi(\cdot\mid s) a ∼ π ( ⋅ ∣ s )
( s ′ , r , d o n e ) ← S T E P ( s , a ) (s', r, \mathrm{done}) \gets \mathrm{STEP}(s, a) ( s ′ , r , done ) ← STEP ( s , a )
If d o n e \mathrm{done} done
Else
y ← r + γ ⋅ V ( s ′ ) y \gets r + \gamma \cdot V(s') y ← r + γ ⋅ V ( s ′ )
V ( s ) ← V ( s ) + α ⋅ ( y − V ( s ) ) V(s)\gets V(s) + \alpha\cdot(y-V(s)) V ( s ) ← V ( s ) + α ⋅ ( y − V ( s ))
s ← s ′ s\gets s' s ← s ′
return V V V
接下来我们来分析一下算法的正确性。
首先,V ‾ t \overline{V}_t V t 被称为 TD target 的原因为
∣ V t + 1 ( s t ) − V ‾ t ( s t ) ∣ = ∣ V t ( s t ) − α t ( s t ) δ t − V ‾ t ( s t ) ∣ = ∣ 1 − α t ( s t ) ∣ ∣ V t ( s t ) − V ‾ t ( s t ) ∣ < ∣ V t ( s t ) − V ‾ t ( s t ) ∣ \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} ∣ V t + 1 ( s t ) − V t ( s t ) ∣ = ∣ V t ( s t ) − α t ( s t ) δ t − V t ( s t ) ∣ = ∣1 − α t ( s t ) ∣∣ V t ( s t ) − V t ( s t ) ∣ < ∣ V t ( s t ) − V t ( s t ) ∣
可以看到函数序列 { V t } \{V_t\} { V t } 收敛到 V ‾ t \overline{V}_t V t .
另一方面,我们有
E [ δ t ∣ S t = s t ] = E [ V π ( S t ) − ( r t + γ V π ( s t + 1 ) ) ∣ S t = s t ] = V π ( s t ) − E [ r t + γ V π ( s t + 1 ) ∣ S t = s t ] = 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} E [ δ t ∣ S t = s t ] = E [ V π ( S t ) − ( r t + γ V π ( s t + 1 ) ) ∣ S t = s t ] = V π ( s t ) − E [ r t + γ V π ( s t + 1 ) ∣ S t = s t ] = 0
因此 TD error δ t \delta_t δ t 衡量了当前的估计 V t V_t V t 与 V π V^\pi V π 之间的差距。
算法的正确性由如下定理给定
Theorem: Convergence of TD Learning 给定策略 π \pi π , 基于 TD learning algorithm , 对任意 s ∈ S s\in\mathcal{S} s ∈ S , 如果 ∑ t α t ( s ) = ∞ \sum_t\alpha_t(s)=\infty ∑ t α t ( s ) = ∞ , ∑ t α t 2 ( s ) < ∞ \sum_t\alpha_t^2(s)<\infty ∑ t α t 2 ( s ) < ∞ , 则 V t ( s ) V_t(s) V t ( s ) a . s . a.s. a . s . 收敛到 V π ( s ) V^{\pi}(s) V π ( s ) , t → ∞ t\to\infty t → ∞ .
Proof:
我们在不同 vocab size 下对比 matrix solve, MC 和 TD.
对于 MC 和 TD, 我们将 episodes 设置为 10 , 000 10,000 10 , 000 , TD 的 learning rate 设置为 0.1 0.1 0.1 .
结果如下图所示
可以看到,matrix exact/iterative solve 每次都需要枚举状态,而 MC/TD 的单条数据成本主要由
rollout 长度决定,因此我们可以在不枚举整个状态空间的情况下更新被访问的状态。
不过,固定 rollout budget 并不保证所有状态都得到充分访问,因此随着状态空间的增大,MC/TD 方法的误差也逐渐增高。
我们采样成功的概率为
p s u c c e s s = ∏ t = 0 m − 1 π ( a t ∗ ∣ a < t ∗ ) . p_{\mathrm{success}}=\prod_{t=0}^{m-1}\pi(a_t^*\mid a_{<t}^*). p success = t = 0 ∏ m − 1 π ( a t ∗ ∣ a < t ∗ ) .
如果 policy 接近均匀分布,该概率大致按 ∣ V ∣ − m |\mathcal V|^{-m} ∣ V ∣ − m 下降。
因此,我们额外统计不同 vocab size 下,MC 和 TD 达到指定误差所需的 episodes.
我们将 maximum episodes 设置为 100 , 000 100,000 100 , 000 ,目标 error 设置为 0.01 0.01 0.01 ,结果如下:
可以看到,在当前 terminal binary reward、固定超参数和误差指标下,MC 比 TD 更快达到目标误差。
当 vocab size 超过 30 后,两者在 rollout budget 内都很难观察到足够多的正样本。
这里暴露出两个不同的问题:
tabular learning 方法不能在相似 prefix 之间共享信息
稀疏 verifier reward 导致有效样本极少。
函数逼近主要缓解前一个问题;后一个问题还需要更好的 exploration, reward shaping 或更强的初始 policy.
上面我们介绍了 MC 和 TD 两种 value function 估计方法,但是当动作空间和状态空间特别大时,tabular learning 仍然存在扩展性问题。
为了解决这个扩展性不足的问题,当状态空间 S \mathcal{S} S 较大或连续时,我们可以使用函数逼近的方法来近似 V π V^{\pi} V π .
这里我们使用 V ϕ : S → R V_{\phi}:\mathcal{S}\to\mathbb{R} V ϕ : S → R 近似 V π V^{\pi} V π ,其中 ϕ \phi ϕ 是参数。
令 d π d^\pi d π 表示策略 π \pi π 在 rollout 中诱导的状态访问分布,我们使用 mean squared value (MSE) error:
L ( ϕ ) = E s ∼ d π [ 1 2 ( 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]. L ( ϕ ) = E s ∼ d π [ 2 1 ( V ϕ ( s ) − V π ( s ) ) 2 ] .
V π ( s ) V^\pi(s) V π ( s ) 满足 V π ( s ) = E π [ G ∣ S = s ] V^\pi(s)=\mathbb E^\pi[G\mid S=s] V π ( s ) = E π [ G ∣ S = s ] , 因此:
∇ ϕ L ( ϕ ) = E s ∼ d π [ ( V ϕ ( s ) − V π ( s ) ) ∇ ϕ V ϕ ( s ) ] = E s ∼ d π , 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} ∇ ϕ L ( ϕ ) = E s ∼ d π [ ( V ϕ ( s ) − V π ( s ) ) ∇ ϕ V ϕ ( s ) ] = E s ∼ d π , G ∼ π ( ⋅ ∣ s ) [ ( V ϕ ( s ) − G ) ∇ ϕ V ϕ ( s ) ] .
现在我们可以使用 MC 方法来估计梯度:
∇ ϕ L ( ϕ ) ≈ g : = 1 M ∑ i = 1 M ( 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)}). ∇ ϕ L ( ϕ ) ≈ g := M 1 i = 1 ∑ M ( V ϕ ( s ( i ) ) − G ( i ) ) ∇ ϕ V ϕ ( s ( i ) ) .
结合 MC 和 SGD 的算法如下:
Algorithm: Value function approximation with MC while not converged:
s 0 ∼ p 0 s_0 \sim p_0 s 0 ∼ p 0 , t = 0 t = 0 t = 0 , τ = [ ] \tau=[] τ = [ ] .
while s t ≠ ⟨ term ⟩ s_t \neq \langle\text{term}\rangle s t = ⟨ term ⟩ :
a t ∼ π ( ⋅ ∣ s t ) a_t \sim \pi(\cdot \mid s_t) a t ∼ π ( ⋅ ∣ s t )
( r t , s t + 1 ) ∼ p ( ⋅ , ⋅ ∣ s t , a t ) (r_t, s_{t+1}) \sim p(\cdot, \cdot \mid s_t, a_t) ( r t , s t + 1 ) ∼ p ( ⋅ , ⋅ ∣ s t , a t )
τ . a p p e n d ( ( s t , r t ) ) \tau.\mathrm{append}((s_t,r_t)) τ . append (( s t , r t ))
t ← t + 1 t \gets t + 1 t ← t + 1
Set T = t T = t T = t .
For every t = 0 , … , T − 1 t=0,\ldots,T-1 t = 0 , … , T − 1 , compute
G t = ∑ k = t T − 1 γ k − t r k G_t=\sum_{k=t}^{T-1}\gamma^{k-t}r_k G t = ∑ k = t T − 1 γ k − t r k .
g ← 1 T ∑ t = 0 T − 1 ( V ϕ ( s t ) − G t ) ∇ ϕ V ϕ ( s t ) g\gets\frac{1}{T}\sum_{t=0}^{T-1}
(V_\phi(s_t)-G_t)\nabla_\phi V_\phi(s_t) g ← T 1 ∑ t = 0 T − 1 ( V ϕ ( s t ) − G t ) ∇ ϕ V ϕ ( s t ) .
Update ϕ \phi ϕ using g g g with an optimizer.
如果能够访问真实的 V π ( s ′ ) V^\pi(s') V π ( s ′ ) ,我们可以将一步使用 TD target 当作监督信号:
L ( ϕ ) = E s ∼ d π π [ 1 2 ( 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 ( ϕ ) = E s ∼ d π π [ 2 1 ( V ϕ ( s ) − ( r + γ V π ( s ′ ) ) ) 2 ] .
对应的梯度为:
∇ ϕ L ( ϕ ) = E s ∼ d π π [ ( 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]. ∇ ϕ L ( ϕ ) = E s ∼ d π π [ ( V ϕ ( s ) − r − γ V π ( s ′ ) ) ∇ ϕ V ϕ ( s ) ] .
这里梯度包含了一个未知量 V π ( s ′ ) V^{\pi}(s') V π ( s ′ ) .
为了解决这个问题,一个自然的想法是使用当前的 value function V ϕ ( s ′ ) V_{\phi}(s') V ϕ ( s ′ ) 来进行代替:
g : = E s ∼ d π π [ ( 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]. g := E s ∼ d π π [ ( V ϕ ( s ) − r − γ V ϕ ( s ′ ) ) ∇ ϕ V ϕ ( s ) ] .
对应的数学更新为:
Algorithm: Value function approximation with TD (not correct) While not converged:
s 0 ∼ p 0 s_0 \sim p_0 s 0 ∼ p 0 , a 0 ∼ π ( ⋅ ∣ s 0 ) a_0 \sim \pi(\cdot \mid s_0) a 0 ∼ π ( ⋅ ∣ s 0 ) , ( r 0 , s 1 ) ∼ p ( ⋅ , ⋅ ∣ s 0 , a 0 ) (r_0, s_1) \sim p(\cdot, \cdot \mid s_0, a_0) ( r 0 , s 1 ) ∼ p ( ⋅ , ⋅ ∣ s 0 , a 0 ) .
compute TD target
y = { r 0 + γ V ϕ ( s 1 ) , if s 1 ≠ ⟨ term ⟩ r 0 , otherwise y = \begin{cases}
r_0 + \gamma V_{\phi}(s_1), & \text{if } s_1 \neq \langle\text{term}\rangle \\
r_0, & \text{otherwise}
\end{cases} y = { r 0 + γ V ϕ ( s 1 ) , r 0 , if s 1 = ⟨ term ⟩ otherwise
g = ( V ϕ ( s 0 ) − y ) ∇ ϕ V ϕ ( s 0 ) g = (V_{\phi}(s_0) - y) \nabla_{\phi} V_{\phi}(s_0) g = ( V ϕ ( s 0 ) − y ) ∇ ϕ V ϕ ( s 0 ) .
Update ϕ \phi ϕ using g g g 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()
可以看到,我们实际上求的梯度是:
∇ ϕ [ 1 2 ( 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} ∇ ϕ [ 2 1 ( V ϕ ( s ) − ( r + γ V ϕ ( s ′ )) ) 2 ] = ( V ϕ ( s ) − ( r + γ V ϕ ( s ′ ))) ( ∇ ϕ V ϕ ( s ) − γ ∇ ϕ V ϕ ( s ′ )) = g
为了解决这个问题,我们可以使用 stop-gradient 技巧来避免 V ϕ ( s ′ ) V_{\phi}(s') V ϕ ( s ′ ) 参与反向传播,此时目标函数变为:
L ( ϕ ) = 1 2 ( V ϕ ( s ) − ( r + γ s g [ V ϕ ( s ′ ) ] ) ) 2 \mathcal{L}(\phi)=\frac{1}{2}\left(V_{\phi}(s) - (r + \gamma \, \mathrm{sg}[V_{\phi}(s')])\right)^2 L ( ϕ ) = 2 1 ( V ϕ ( s ) − ( r + γ sg [ V ϕ ( s ′ )]) ) 2
其中 s g [ ⋅ ] \mathrm{sg}[\cdot] sg [ ⋅ ] 是 stop-gradient operator, 满足:
s g [ x ] = { x forward pass 0 backward pass \mathrm{sg}[x] = \begin{cases}
x & \text{forward pass} \\
0 & \text{backward pass}
\end{cases} sg [ x ] = { x 0 forward pass backward pass
对应的梯度就是:
∇ ϕ [ 1 2 ( V ϕ ( s ) − ( r + γ s g [ 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 ∇ ϕ [ 2 1 ( V ϕ ( s ) − ( r + γ sg [ V ϕ ( s ′ )]) ) 2 ] = ( V ϕ ( s ) − ( r + γ V ϕ ( s ′ ))) ∇ ϕ V ϕ ( 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:
s 0 ∼ p 0 s_0 \sim p_0 s 0 ∼ p 0 , a 0 ∼ π ( ⋅ ∣ s 0 ) a_0 \sim \pi(\cdot \mid s_0) a 0 ∼ π ( ⋅ ∣ s 0 ) , ( r 0 , s 1 ) ∼ p ( ⋅ , ⋅ ∣ s 0 , a 0 ) (r_0, s_1) \sim p(\cdot, \cdot \mid s_0, a_0) ( r 0 , s 1 ) ∼ p ( ⋅ , ⋅ ∣ s 0 , a 0 ) .
compute TD target
y = { r 0 + γ s g [ V ϕ ( s 1 ) ] , if s 1 ≠ ⟨ term ⟩ r 0 , otherwise y = \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} y = { r 0 + γ sg [ V ϕ ( s 1 )] , r 0 , if s 1 = ⟨ term ⟩ otherwise
g = ( V ϕ ( s 0 ) − y ) ∇ ϕ V ϕ ( s 0 ) g = (V_{\phi}(s_0) - y) \nabla_{\phi} V_{\phi}(s_0) g = ( V ϕ ( s 0 ) − y ) ∇ ϕ V ϕ ( s 0 ) .
Update ϕ \phi ϕ using g g g with an optimizer.
实际上,我们可以进一步推广 TD 到多步场景,注意到:
V π ( s 0 ) = E π [ r 0 + γ V π ( s 1 ) ∣ s 0 ] = E π [ r 0 + γ E π [ r 1 + γ V π ( s 2 ) ∣ s 1 ] ∣ s 0 ] = E π [ r 0 + γ r 1 + γ 2 V π ( s 2 ) ∣ s 0 ] . \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 π ( s 0 ) = E π [ r 0 + γ V π ( s 1 ) ∣ s 0 ] = E π [ r 0 + γ E π [ r 1 + γ V π ( s 2 ) ∣ s 1 ] ∣ s 0 ] = E π [ r 0 + γ r 1 + γ 2 V π ( s 2 ) ∣ s 0 ] .
因此:
V π ( s 0 ) = E π [ r 0 + γ r 1 + γ 2 V π ( s 2 ) ∣ s 0 ] V^{\pi}(s_0) = \mathbb{E}^{\pi}[r_0 + \gamma r_1 + \gamma^2 V^{\pi}(s_2) \mid s_0] V π ( s 0 ) = E π [ r 0 + γ r 1 + γ 2 V π ( s 2 ) ∣ s 0 ]
重复这个过程 k k k 次,我们就可以得到 k k k step TD:
V π ( s 0 ) = E π [ ∑ i = 0 k − 1 γ i r i + γ k V π ( s k ) ∣ s 0 ] 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] V π ( s 0 ) = E π [ i = 0 ∑ k − 1 γ i r i + γ k V π ( s k ) ∣ s 0 ]
基于 k-step transition ,我们可以构建对应的目标函数并应用 stop-gradient:
Algorithm: Value function approximation with k-step TD while not converged:
s 0 ∼ p 0 s_0 \sim p_0 s 0 ∼ p 0 , set t = 0 t = 0 t = 0 .
for t = 0 , … , k − 1 t = 0, \ldots, k-1 t = 0 , … , k − 1 , stopping early if the episode terminates:
a t ∼ π ( ⋅ ∣ s t ) a_t \sim \pi(\cdot \mid s_t) a t ∼ π ( ⋅ ∣ s t )
( r t , s t + 1 ) ∼ p ( ⋅ , ⋅ ∣ s t , a t ) (r_t, s_{t+1}) \sim p(\cdot, \cdot \mid s_t, a_t) ( r t , s t + 1 ) ∼ p ( ⋅ , ⋅ ∣ s t , a t )
Let n ≤ k n\leq k n ≤ k be the number of collected transitions and compute
y = { ∑ t = 0 n − 1 γ t r t + γ n s g [ V ϕ ( s n ) ] , if s n ≠ ⟨ term ⟩ , ∑ t = 0 n − 1 γ t r t , 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} y = { ∑ t = 0 n − 1 γ t r t + γ n sg [ V ϕ ( s n )] , ∑ t = 0 n − 1 γ t r t , if s n = ⟨ term ⟩ , otherwise .
g ← ( V ϕ ( s 0 ) − y ) ∇ ϕ V ϕ ( s 0 ) g \gets (V_{\phi}(s_0) - y) \nabla_{\phi} V_{\phi}(s_0) g ← ( V ϕ ( s 0 ) − y ) ∇ ϕ V ϕ ( s 0 ) .
Update ϕ \phi ϕ using g g g with an optimizer.
对于 Q-function, 我们也可以和 value function 一样设计类似的算法。对于 MC, 有:
Q π ( s 0 , a 0 ) = E π [ ∑ t = 0 T − 1 γ t r t ∣ s 0 , a 0 ] ≈ 1 M ∑ i = 1 M ∑ t = 0 T ( i ) − 1 γ t r t ( 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)}. Q π ( s 0 , a 0 ) = E π [ t = 0 ∑ T − 1 γ t r t ∣ s 0 , a 0 ] ≈ M 1 i = 1 ∑ M t = 0 ∑ T ( i ) − 1 γ t r t ( i ) .
当使用函数来近似时,目标函数为:
L ( ϕ ) = E s ∼ d π , a ∼ π ( ⋅ ∣ s ) [ 1 2 ( 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 s ∼ d π , a ∼ π ( ⋅ ∣ s ) [ 2 1 ( Q ϕ ( s , a ) − Q π ( s , a ) ) 2 ] .
对应的梯度:
∇ ϕ L ( ϕ ) = E π [ ( Q ϕ ( s t , a t ) − G t ) ∇ ϕ Q ϕ ( s t , a t ) ] . \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]. ∇ ϕ L ( ϕ ) = E π [ ( Q ϕ ( s t , a t ) − G t ) ∇ ϕ Q ϕ ( s t , a t ) ] .
对应的算法如下
Algorithm: Q-function approximation with MC while not converged:
s 0 ∼ p 0 s_0 \sim p_0 s 0 ∼ p 0 , set t = 0 t = 0 t = 0 , τ = [ ] \tau=[] τ = [ ] .
while s t ≠ ⟨ term ⟩ s_t \neq \langle\text{term}\rangle s t = ⟨ term ⟩ :
a t ∼ π ( ⋅ ∣ s t ) a_t \sim \pi(\cdot \mid s_t) a t ∼ π ( ⋅ ∣ s t )
( r t , s t + 1 ) ∼ p ( ⋅ , ⋅ ∣ s t , a t ) (r_t, s_{t+1}) \sim p(\cdot, \cdot \mid s_t, a_t) ( r t , s t + 1 ) ∼ p ( ⋅ , ⋅ ∣ s t , a t )
τ . a p p e n d ( ( s t , a t , r t ) ) \tau.\mathrm{append}((s_t,a_t,r_t)) τ . append (( s t , a t , r t ))
t ← t + 1 t \gets t + 1 t ← t + 1
Set T = t T = t T = t .
For every t t t , compute G t = ∑ k = t T − 1 γ k − t r k G_t=\sum_{k=t}^{T-1}\gamma^{k-t}r_k G t = ∑ k = t T − 1 γ k − t r k .
g = 1 T ∑ t = 0 T − 1 ( Q ϕ ( s t , a t ) − G t ) ∇ ϕ Q ϕ ( s t , a t ) 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) g = T 1 ∑ t = 0 T − 1 ( Q ϕ ( s t , a t ) − G t ) ∇ ϕ Q ϕ ( s t , a t ) .
Update ϕ \phi ϕ using g g g with an optimizer.
Return Q ϕ Q_\phi Q ϕ
对于 TD learning, 我们的目标函数为
L ( ϕ ) = E s ∼ d π , a ∼ π ( ⋅ ∣ s ) [ 1 2 ( 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], L ( ϕ ) = E s ∼ d π , a ∼ π ( ⋅ ∣ s ) [ 2 1 ( Q ϕ ( s , a ) − r − γ Q π ( s ′ , a ′ ) ) 2 ] ,
其中非终止状态下 a ′ ∼ π ( ⋅ ∣ s ′ ) a'\sim\pi(\cdot\mid s') a ′ ∼ π ( ⋅ ∣ s ′ ) .
与 value function 类似,使用 bootstrap target 可以得到 on-policy 的 Expected SARSA/SARSA 风格 policy evaluation:
Algorithm: Q-function approximation with TD While not converged:
s 0 ∼ p 0 s_0 \sim p_0 s 0 ∼ p 0 , a 0 ∼ π ( ⋅ ∣ s 0 ) a_0 \sim \pi(\cdot \mid s_0) a 0 ∼ π ( ⋅ ∣ s 0 ) , ( r 0 , s 1 ) ∼ p ( ⋅ , ⋅ ∣ s 0 , a 0 ) (r_0, s_1) \sim p(\cdot, \cdot \mid s_0, a_0) ( r 0 , s 1 ) ∼ p ( ⋅ , ⋅ ∣ s 0 , a 0 ) .
y = { r 0 + γ s g [ Q ϕ ( s 1 , a 1 ) ] , a 1 ∼ π ( ⋅ ∣ s 1 ) , if s 1 ≠ ⟨ term ⟩ , r 0 , 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} y = { r 0 + γ sg [ Q ϕ ( s 1 , a 1 )] , a 1 ∼ π ( ⋅ ∣ s 1 ) , r 0 , if s 1 = ⟨ term ⟩ , otherwise .
g = ( Q ϕ ( s 0 , a 0 ) − y ) ∇ ϕ Q ϕ ( s 0 , a 0 ) g = \left(Q_{\phi}(s_0, a_0) - y\right) \nabla_{\phi} Q_{\phi}(s_0, a_0) g = ( Q ϕ ( s 0 , a 0 ) − y ) ∇ ϕ Q ϕ ( s 0 , a 0 ) .
Update ϕ \phi ϕ using g g g with an optimizer.
Return Q ϕ Q_\phi Q ϕ
Algorithm: Q-function approximation with k-step TD While not converged:
s 0 ∼ p 0 s_0 \sim p_0 s 0 ∼ p 0 , set t = 0 t = 0 t = 0 .
For t = 0 , … , k − 1 t = 0, \ldots, k-1 t = 0 , … , k − 1 , stopping early on termination:
a t ∼ π ( ⋅ ∣ s t ) a_t \sim \pi(\cdot \mid s_t) a t ∼ π ( ⋅ ∣ s t )
( r t , s t + 1 ) ∼ p ( ⋅ , ⋅ ∣ s t , a t ) (r_t, s_{t+1}) \sim p(\cdot, \cdot \mid s_t, a_t) ( r t , s t + 1 ) ∼ p ( ⋅ , ⋅ ∣ s t , a t )
Let n ≤ k n\leq k n ≤ k be the number of collected transitions.
If s n s_n s n is non-terminal
sample a n ∼ π ( ⋅ ∣ s n ) a_n\sim\pi(\cdot\mid s_n) a n ∼ π ( ⋅ ∣ s n )
set y = ∑ t = 0 n − 1 γ t r t + γ n s g [ Q ϕ ( s n , a n ) ] y=\sum_{t=0}^{n-1}\gamma^t r_t+\gamma^n\mathrm{sg}[Q_\phi(s_n,a_n)] y = ∑ t = 0 n − 1 γ t r t + γ n sg [ Q ϕ ( s n , a n )] ;
Else
set y = ∑ t = 0 n − 1 γ t r t y=\sum_{t=0}^{n-1}\gamma^t r_t y = ∑ t = 0 n − 1 γ t r t .
g = ( Q ϕ ( s 0 , a 0 ) − y ) ∇ ϕ Q ϕ ( s 0 , a 0 ) g = (Q_{\phi}(s_0, a_0) - y) \nabla_{\phi} Q_{\phi}(s_0, a_0) g = ( Q ϕ ( s 0 , a 0 ) − y ) ∇ ϕ Q ϕ ( s 0 , a 0 )
Update ϕ \phi ϕ using g g g with an optimizer.
Return Q ϕ Q_\phi Q ϕ
介绍完 policy evaluation 之后,我们可以设计 policy iteration 算法,其基本思想是交替进行 policy evaluation 和 policy improvement:
Compute V π k and Q π k (Policy evaluation) π k + 1 ( s ) = arg max a ∈ A E ( 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} Compute V π k and Q π k π k + 1 ( s ) = arg a ∈ A max E ( r , s ′ ) ∼ p ( ⋅ , ⋅ ∣ s , a ) [ r + γ V π k ( s ′ ) ∣ s , a ] (Policy evaluation) (Policy improvement)
算法的正确性由如下定理保证
Theorem: Theorem : 考虑上述 policy iteration 算法,假设 γ ∈ ( 0 , 1 ) \gamma \in (0, 1) γ ∈ ( 0 , 1 ) , ∣ S ∣ < ∞ |\mathcal{S}| < \infty ∣ S ∣ < ∞ , ∣ A ∣ < ∞ |\mathcal{A}| < \infty ∣ A ∣ < ∞ 以及 ∣ r ∣ ≤ R < ∞ |r| \leq R < \infty ∣ r ∣ ≤ R < ∞ , a.s. 那么:
V π k + 1 ≥ V π k , k = 1 , 2 , … V^{\pi^{k+1}} \geq V^{\pi^k}, \quad k = 1, 2, \dots V π k + 1 ≥ V π k , k = 1 , 2 , … 并且存在 K ∈ N K \in \mathbb{N} K ∈ N 满足 V π K = V ∗ V^{\pi^K} = V^* V π K = V ∗ . 也就是说,有限步之后就可以得到最优策略.
Proof: 由 π k + 1 \pi^{k+1} π k + 1 的 greedy 定义,对每个状态 s s s 都有
( T π k + 1 V π k ) ( s ) = ( T ∗ V π k ) ( s ) ≥ ( T π k V π 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). ( T π k + 1 V π k ) ( s ) = ( T ∗ V π k ) ( s ) ≥ ( T π k V π k ) ( s ) = V π k ( s ) . 反复应用单调的 Bellman operator T π k + 1 \mathcal T^{\pi^{k+1}} T π k + 1 ,并利用其 contraction
性质,得到
V π k + 1 = lim n → ∞ ( T π k + 1 ) n V π k ≥ V π k . V^{\pi^{k+1}}
=\lim_{n\to\infty}
(\mathcal T^{\pi^{k+1}})^nV^{\pi^k}
\geq V^{\pi^k}. V π k + 1 = n → ∞ lim ( T π k + 1 ) n V π k ≥ V π k . 有限 MDP 只有有限个 deterministic policies.
若 value 尚未最优,policy improvement 会产生严格改进;采用一致的 tie-breaking 后不会在等价 policy之间循环。
因此有限次迭代后到达满足 Bellman optimality equation 的 policy, 其 value 即为 V ∗ V^* V ∗ .
MC 必须等待 episode 结束,然后用真实 return 监督每个已访问状态。
对固定状态而言,sample return 是 V π V^\pi V π 的无偏估计,但长轨迹中的 variance 可能很高。
TD 使用 V ϕ ( s t + 1 ) V_\phi(s_{t+1}) V ϕ ( s t + 1 ) bootstrap,不需要等待完整 episode,通常 variance 更低,但 target 会受到当前 value approximation error 的影响。
这里所谓的 “TD 有偏” 主要是指有限训练阶段的 bootstrap target;在满足收敛条件的 tabular on-policy 场景中,TD 仍然可以收敛到真实的 V π V^\pi V π .
k-step TD 在两者之间插值:
TD(0) ← k → 0 k-step TD → k → T MC . \text{TD(0)}
\xleftarrow{\quad k\to0\quad}
\text{k-step TD}
\xrightarrow{\quad k\to T\quad}
\text{MC}. TD(0) k → 0 k-step TD k → T MC .
LLM generation 的 state 是 (prompt, token prefix), 状态空间随 horizon (context length) 指数增长
Matrix solve 和 dynamic programming 能提供小规模环境 ( < 20 < 20 < 20 ) 的精确基线,但无法扩展到真实 LLM 的场景下 (∼ 15 k \sim15k ∼ 15 k )
MC 使用完整 return, low bias, high variance; TD 使用 bootstrap, low variance, high bias; k-step return 在两者之间实现一个 trade off.
Function approximation 让不同 prefix 共享信息,一条 rollout 可以监督多个 prefix
在 binary verifier reward 下,V π ( x , y < t ) V^\pi(x,y_{<t}) V π ( x , y < t ) 可以解释为当前 policy 最终通过 verifier 的概率
Policy evaluation 解决的是“当前 policy 有多好”;下一步 policy improvement 才解决“如何让 policy 变得更好”
在 Policy Evaluation 一节中,策略 π \pi π 是固定的,我们只估计 V π V^\pi V π 或 Q π Q^\pi Q π .
Value-based methods 更进一步:学习能够区分动作优劣的 value function,并据此改进策略。
一种直接思路是寻找 optimal value function V ∗ V^* V ∗ 或 optimal Q function Q ∗ Q^* Q ∗ ,再从中恢复最优策略。
这也是 value based methods 的核心思想:
求解得到 optimal value function V ∗ V^* V ∗ 或者 optimal Q function Q ∗ Q^* Q ∗ , 再得到 optimal policy π ∗ \pi^* π ∗ .
Value iteration 的基本思路为:通过 Bellman optimality operator 迭代求解 V ∗ V^* V ∗ 和 Q ∗ Q^* Q ∗ , 然后得到最优策略。
根据 BOE Theorem 的证明,T ∗ \mathcal{T}^* T ∗ 是一个 contraction mapping, 因此我们可以构造函数序列:
V , T ∗ V , ( T ∗ ) 2 V , ⋯ V, \mathcal{T}^* V, (\mathcal{T}^*)^2 V, \cdots V , T ∗ V , ( T ∗ ) 2 V , ⋯
由不动点定理,该序列收敛到 V ∗ V^* V ∗ .
这就是 value iteration algorithm:
V k + 1 = T ∗ V k V k + 1 ( s t ) = max a ∈ A E ( r t , s t + 1 ) ∼ p ( ⋅ , ⋅ ∣ s t , a t ) [ r + γ V k ( s t + 1 ) ∣ S t = s t , A t = a t ] \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} V k + 1 V k + 1 ( s t ) = T ∗ V k = a ∈ A max E ( r t , s t + 1 ) ∼ p ( ⋅ , ⋅ ∣ s t , a t ) [ r + γ V k ( s t + 1 ) ∣ S t = s t , A t = a t ]
类似地,我们还有 Q value iteration algorithm:
Q k + 1 = T ∗ Q k Q k + 1 ( s t , a t ) = E ( r t , s t + 1 ) ∼ p ( ⋅ , ⋅ ∣ s t , a t ) [ r + γ max a t + 1 ∈ A Q k ( s t + 1 , a t + 1 ) ∣ S t = s t , A t = a t ] \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} Q k + 1 Q k + 1 ( s t , a t ) = T ∗ Q k = E ( r t , s t + 1 ) ∼ p ( ⋅ , ⋅ ∣ s t , a t ) [ r + γ a t + 1 ∈ A max Q k ( s t + 1 , a t + 1 ) ∣ S t = s t , A t = a t ]
最终,我们的最优策略为
π ∗ ( s t ) = arg max a ∈ A Q ∗ ( s t , a ) \pi^{*}(s_t) = \arg\max_{a \in \mathcal{A}} Q^*(s_t,a) π ∗ ( s t ) = arg a ∈ A max Q ∗ ( s t , a )
其中 Q ∗ = lim k → ∞ Q k Q^*=\lim_{k\to\infty}Q^k Q ∗ = lim k → ∞ Q k .
本节将 value iteration 的思想与函数逼近、TD learning 等技术结合,得到实用的 value-based 算法.
当 state 和 action space 足够小时,我们可以直接维护一张 Q-table.
在 Model-Free 设定下,我们无法直接计算转移概率 E s ′ ∼ p \mathbb{E}_{s' \sim p} E s ′ ∼ p .
此时,我们必须通过与环境交互产生的轨迹样本 ( s , a , r , s ′ , a ′ ) (s, a, r, s', a') ( s , a , r , s ′ , a ′ ) 来更新价值。
SARSA 的核心在于它的更新动作 a ′ a' a ′ 严格来自于当前的 ϵ \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 ( s , a ) ← Q ( s , a ) + α [ r + γ Q ( s ′ , a ′ ) − Q ( s , a ) ]
与 SARSA 不同,Q-Learning 在估计未来价值时,直接贪心地选择最优动作,而不受当前探索策略的影响:
Q ( s , a ) ← Q ( s , a ) + α [ r + γ max a ′ ∈ A Q ( 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]} Q ( s , a ) ← Q ( s , a ) + α [ r + γ a ′ ∈ A max Q ( s ′ , a ′ ) − Q ( s , a ) ]
当状态空间较大或连续时,我们需要使用函数逼近,与上一节类似,我们用 V ϕ V_{\phi} V ϕ 近似 V π V^{\pi} V π :
J ( ϕ ) = E s [ 1 2 ( 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 ( ϕ ) = E s [ 2 1 ( V π ( s ) − V ϕ ( s ) ) 2 ]
目标函数的梯度为:
∇ ϕ J ( ϕ ) = − E s [ ( 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] ∇ ϕ J ( ϕ ) = − E s [ ( V π ( s ) − V ϕ ( s )) ∇ ϕ V ϕ ( s ) ]
随机梯度下降的更新形式为:
ϕ 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) ϕ k + 1 = ϕ k + α k ( V π ( s ) − V ϕ ( s )) ∇ ϕ V ϕ ( s )
V π ( s ) V^{\pi}(s) V π ( s ) 可以通过两种方式来近似:
MC 近似 : V π ( s ) = E π [ G 0 ∣ s 0 = s ] ≈ G 0 V^{\pi}(s) = \mathbb{E}^{\pi}[G_0 \mid s_0 = s] \approx G_0 V π ( s ) = E π [ G 0 ∣ s 0 = s ] ≈ G 0
TD 近似 : V π ( s ) ≈ r t + γ V ϕ ( s t + 1 ) V^{\pi}(s) \approx r_t + \gamma V_{\phi}(s_{t+1}) V π ( s ) ≈ r t + γ V ϕ ( s t + 1 )
将 function approximation 与 Q-learning 结合,可以用 neural network
逼近 optimal Q function。DQN 进一步引入 replay buffer 和 target network
来改善训练稳定性:
L ( ϕ ) = E ( s , a , r , s ′ ) ∼ D [ 1 2 ( r + γ max a ′ Q ϕ − ( 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] L ( ϕ ) = E ( s , a , r , s ′ ) ∼ D [ 2 1 ( r + γ a ′ max Q ϕ − ( s ′ , a ′ ) − Q ϕ ( s , a ) ) 2 ]
其中 ϕ − \phi^- ϕ − 为**目标网络(Target Network)**的参数,用于切断数据自举(Bootstrapping)带来的训练正反馈爆炸。
为了保证求导时 Target 的稳定性,我们必须在代码或公式里明确使用 stop-gradient 算子:
y = r + γ max a ′ s g [ Q ϕ − ( s ′ , a ′ ) ] y = r + \gamma \max_{a'} \mathrm{sg}\left[Q_{\phi^-}(s', a')\right] y = r + γ max a ′ sg [ Q ϕ − ( s ′ , a ′ ) ]
DQN 针对 optimal policy 应用了 TD learning, 其目标函数为:
min ϕ L ( ϕ ) = E ( s , a , r , s ′ ) ∼ D [ 1 2 ( Q ϕ ( s , a ) − ( r + γ max a ′ Q ϕ ( 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] ϕ min L ( ϕ ) = E ( s , a , r , s ′ ) ∼ D [ 2 1 ( Q ϕ ( s , a ) − ( r + γ a ′ max Q ϕ ( s ′ , a ′ ) ) ) 2 ]
这里 D \mathcal{D} D 是 replay buffer, 存储了过去的 transition ( s , a , r , s ′ ) (s, a, r, s') ( s , a , r , s ′ ) .
💡
DQN 的两个关键技术:
Experience Replay : 使用 replay buffer 存储过去的 transitions,从中随机采样进行训练。这打破了数据之间的相关性,提高了样本效率.
Target Network : 使用一个独立的 target network Q ϕ − Q_{\phi^{-}} Q ϕ − 来计算 TD target, 其参数 ϕ − \phi^{-} ϕ − 定期从 ϕ \phi ϕ 复制(或通过 Polyak averaging 更新)。这稳定了训练过程.
带有 target network 和 stop-gradient 的 DQN 目标函数为:
L ( ϕ ) = E ( s , a , r , s ′ ) ∼ D [ 1 2 ( Q ϕ ( s , a ) − ( r + γ max a ′ s g [ 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] L ( ϕ ) = E ( s , a , r , s ′ ) ∼ D [ 2 1 ( Q ϕ ( s , a ) − ( r + γ a ′ max sg [ Q ϕ − ( s ′ , a ′ )] ) ) 2 ]
对应的 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
Initialize replay buffer D \mathcal{D} D with capacity N N N .
Initialize Q-network Q ϕ Q_{\phi} Q ϕ with random weights.
Initialize target network Q ϕ − Q_{\phi^{-}} Q ϕ − with weights ϕ − = ϕ \phi^{-} = \phi ϕ − = ϕ .
for episode = 1 , … , M 1, \dots, M 1 , … , M :
s 0 ∼ p 0 s_0 \sim p_0 s 0 ∼ p 0 .
for t = 0 , … , T − 1 t = 0, \dots, T-1 t = 0 , … , T − 1 :
Select action with ϵ \epsilon ϵ -greedy:
a t = { arg max a Q ϕ ( s t , 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} a t = { arg max a Q ϕ ( s t , a ) , random action , with prob 1 − ϵ with prob ϵ
( r t , s t + 1 ) ∼ p ( ⋅ , ⋅ ∣ s t , a t ) (r_t, s_{t+1}) \sim p(\cdot, \cdot \mid s_t, a_t) ( r t , s t + 1 ) ∼ p ( ⋅ , ⋅ ∣ s t , a t ) .
Store ( s t , a t , r t , s t + 1 ) (s_t, a_t, r_t, s_{t+1}) ( s t , a t , r t , s t + 1 ) in D \mathcal{D} D .
Sample random minibatch { ( s i , a i , r i , s i ′ ) } \{(s_i, a_i, r_i, s'_i)\} {( s i , a i , r i , s i ′ )} from D \mathcal{D} D .
Compute target:
y i = { r i + γ max a ′ Q ϕ − ( s i ′ , a ′ ) , if not terminal r i , if terminal y_i = \begin{cases}
r_i + \gamma \max_{a'} Q_{\phi^{-}}(s'_i, a'), & \text{if not terminal} \\
r_i, & \text{if terminal}
\end{cases} y i = { r i + γ max a ′ Q ϕ − ( s i ′ , a ′ ) , r i , if not terminal if terminal
Update Q-network:
L = 1 ∣ batch ∣ ∑ i 1 2 ( Q ϕ ( s i , a i ) − y i ) 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} L = ∣ batch ∣ 1 i ∑ 2 1 ( Q ϕ ( s i , a i ) − y i ) 2 , ϕ ← ϕ − α ∇ ϕ L
Every C C C steps: ϕ − ← ϕ \phi^{-} \gets \phi ϕ − ← ϕ .
Value-based control 建立了从 Bellman optimality equation 到可训练算法的桥梁,
但 DQN 通常不是 LLM post-training 的最终选择:
LLM action space 是整个 vocabulary,每个 prefix 都要对大量 token 计算并比较 Q ( s , a ) Q(s,a) Q ( s , a )
arg max a Q ( s , a ) \arg\max_a Q(s,a) arg max a Q ( s , a ) 产生的 greedy policy 不等于我们希望保留和微调的 autoregressive policy distribution
LLM reward 常在完整 response 后给出,早期 token 的 Q target 具有严重的 credit-assignment 和 exploration 问题
预训练 LLM 已经提供了强 policy prior,直接优化 π θ \pi_\theta π θ 通常比从 Q function 重新构造生成策略更自然
因此,本章的主要作用是解释 evaluation 如何变成 control , 并为 TD, off-policy learning 和 bootstrap 建立直觉。
下一章将转向 policy-based methods:直接对 LLM policy 的 expected return 求梯度。
Value-based methods 先学习 Q ∗ Q^* Q ∗ ,再通过 arg max a Q ∗ ( s , a ) \arg\max_a Q^*(s,a) arg max a Q ∗ ( s , a ) 间接得到策略。
对 LLM 而言,我们已经有一个参数化的 autoregressive policy π θ ( y t ∣ x , y < t ) \pi_\theta(y_t\mid x,y_{<t}) π θ ( y t ∣ x , y < t ) , 因此更自然的选择是直接调整 θ \theta θ ,提高高回报 response 的生成概率。
这就是 policy-based methods 的核心思想。
使用 policy π θ \pi_\theta π θ 生成 trajectory, 通过 trajectory reward 构造 gradient estimator,再直接改进当前策略 π θ \pi_\theta π θ .
根据采样策略 (behavior policy) 的不同,policy-based methods 一般会被分为两类:
on-policy: 采集数据的 policy 和训练更新的 policy 一致
off-policy: 采集数据的 policy 和训练更新的 policy 不一致
这两者的核心区别在于:我们更新模型所使用的数据是否由当前模型产生? 如果是,则说明算法是 On-policy, 反之则说明是 Off-policy.
on-policy 和 off-policy 对比如下表所示
目前,主流针对 LLM 的 RL 算法基本都是
我们的目标与前面一致,仍然是最大化 return , 问题形式为
max π J ( θ ) = E τ ∼ π θ [ G 0 ( τ ) ] \max_{\pi}\mathcal{J}(\theta)=\mathbb{E}_{\tau\sim\pi_\theta}\left[G_0(\tau)\right] π max J ( θ ) = E τ ∼ π θ [ G 0 ( τ ) ]
我们先来求解目标函数的梯度,梯度由下面的 policy gradient theorem 给出。
Theorem: Policy Gradient Theorem 目标函数 最大化 return 的梯度为
∇ θ J ( θ ) = E τ ∼ π θ [ ∑ t = 0 T − 1 ∇ θ log π θ ( a t ∣ s t ) G 0 ( τ ) ] \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] ∇ θ J ( θ ) = E τ ∼ π θ [ t = 0 ∑ T − 1 ∇ θ log π θ ( a t ∣ s t ) G 0 ( τ ) ]
Proof: 由 trajectory definition , trajectory τ \tau τ 的概率分布可以写为
p θ ( τ ) = p 0 ( s 0 ) ∏ t = 0 T − 1 p ( r t , s t + 1 ∣ s t , a t ) π θ ( a t ∣ s t ) 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 0 ( s 0 ) t = 0 ∏ T − 1 p ( r t , s t + 1 ∣ s t , a t ) π θ ( a t ∣ s t ) 对 p θ ( τ ) p_\theta(\tau) p θ ( τ ) 求梯度得到
∇ θ p θ ( τ ) = p θ ( τ ) ∇ θ log p θ ( τ ) = p θ ( τ ) ∇ θ log [ p 0 ( s 0 ) ∏ t = 0 T − 1 p ( r t , s t + 1 ∣ s t , a t ) π θ ( a t ∣ s t ) ] = p θ ( τ ) ∑ t = 0 T − 1 log π θ ( a t ∣ s t ) \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} ∇ θ p θ ( τ ) = p θ ( τ ) ∇ θ log p θ ( τ ) = p θ ( τ ) ∇ θ log [ p 0 ( s 0 ) t = 0 ∏ T − 1 p ( r t , s t + 1 ∣ s t , a t ) π θ ( a t ∣ s t ) ] = p θ ( τ ) t = 0 ∑ T − 1 log π θ ( a t ∣ s t ) 这样,目标函数的梯度为
∇ θ J ( θ ) = ∇ θ E τ ∼ π θ [ G 0 ( τ ) ] = ∑ τ G 0 ( τ ) ∇ θ p θ ( τ ) = ∑ τ G 0 ( τ ) p θ ( τ ) ∑ t = 0 T − 1 log π θ ( a t ∣ s t ) = E τ ∼ π θ [ ∑ t = 0 T − 1 ∇ θ log π θ ( a t ∣ s t ) G 0 ( τ ) ] \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} ∇ θ J ( θ ) = ∇ θ E τ ∼ π θ [ G 0 ( τ ) ] = τ ∑ G 0 ( τ ) ∇ θ p θ ( τ ) = τ ∑ G 0 ( τ ) p θ ( τ ) t = 0 ∑ T − 1 log π θ ( a t ∣ s t ) = E τ ∼ π θ [ t = 0 ∑ T − 1 ∇ θ log π θ ( a t ∣ s t ) G 0 ( τ ) ] 证毕。
注意 policy gradient theorem 将梯度也写成了期望的形式,这方面我们使用 MC 方法来进行估计和优化。
接下来,我们介绍最简单的 policy gradient methods, 令
g ^ : = ∑ t = 0 T − 1 ∇ θ log π θ ( a t ∣ s t ) G 0 ( τ ) \hat{g} := \sum_{t=0}^{T-1}\nabla_\theta\log\pi_\theta(a_t\mid s_t)G_0(\tau) g ^ := t = 0 ∑ T − 1 ∇ θ log π θ ( a t ∣ s t ) G 0 ( τ )
由 policy gradient theorem, 我们知道 g ^ \hat{g} g ^ 是 ∇ θ J ( θ ) \nabla_\theta \mathcal{J}(\theta) ∇ θ J ( θ ) 的一个无偏估计。
基于 MC 思想,我们可以得到最简单的 policy gradient algorithm, 也就是 REINFORCE 算法
Algorithm: Policy Gradient with MC (REINFORCE) Initialize : policy parameters θ 0 \theta_0 θ 0 .
For k = 0 , 1 , … k=0,1,\dots k = 0 , 1 , … :
Sample M M M trajectories { τ i } \{\tau_i\} { τ i } according to policy π θ k \pi_{\theta_k} π θ k .
Estimate the gradient via MC:
g ^ = 1 M ∑ i = 1 M ∑ t = 0 T i − 1 ∇ θ log π θ k ( a t ∣ s t ) G 0 ( τ 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) g ^ = M 1 i = 1 ∑ M t = 0 ∑ T i − 1 ∇ θ log π θ k ( a t ∣ s t ) G 0 ( τ i )
Update θ k \theta_k θ k using g ^ \hat{g} g ^ with an optimizer.
Return π θ k \pi_{\theta_k} π θ k
虽然 g g g 是 ∇ θ J ( θ ) \nabla_\theta \mathcal{J}(\theta) ∇ θ J ( θ ) 的一个无偏估计,但是使用 g g g 作为估计会出现比较大的方差,TODO
因此我们本节主要介绍如何减少 g g g 作为估计的方差,即 variance reduction.
我们首先介绍 policy gradient 的 baseline invariance 性质。
Proposition: Baseline Invariance of Policy Gradient 针对 policy gradient theorem, 我们有如下等式
E τ ∼ ( p 0 , π θ , p ) [ ∑ t = 0 T − 1 ∇ θ log π θ ( a t ∣ s t ) ∑ t = 0 T − 1 r ( s t ) ] = E τ ∼ ( p 0 , π θ , p ) [ ∑ t = 0 T − 1 ∇ θ log π θ ( a t ∣ s t ) ( ∑ t = 0 T − 1 r ( s t ) − b ( s t ) ) ] \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} = E τ ∼ ( p 0 , π θ , p ) [ t = 0 ∑ T − 1 ∇ θ log π θ ( a t ∣ s t ) t = 0 ∑ T − 1 r ( s t ) ] E τ ∼ ( p 0 , π θ , p ) [ t = 0 ∑ T − 1 ∇ θ log π θ ( a t ∣ s t ) ( t = 0 ∑ T − 1 r ( s t ) − b ( s t ) ) ] 其中 b ( s t ) b(s_t) b ( s t ) 是一个仅与状态有关的函数或者常数,我们将其称为 baseline .
Proof: 我们证明
E τ ∼ ( p 0 , π θ , p ) [ b ( s t ) ∇ θ log p θ ( a t ∣ s t ) ] = 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 τ ∼ ( p 0 , π θ , p ) [ b ( s t ) ∇ θ log p θ ( a t ∣ s t ) ] = 0 注意到
E τ ∼ ( p 0 , π θ , p ) [ b ( s t ) ∇ θ log p θ ( a t ∣ s t ) ] = E s 0 ∼ p 0 [ b ( s t ) E a t ∼ π θ ( ⋅ ∣ s t ) [ ∇ θ log π θ ( a t ∣ s t ) ] ] = E s 0 ∼ p 0 [ b ( s t ) ∑ a ∇ θ π θ ( a ∣ s ) ] = E s 0 ∼ p 0 [ b ( s t ) ∇ θ ∑ a π θ ( a ∣ s ) ] = E s 0 ∼ p 0 [ b ( s t ) ∇ θ 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} E τ ∼ ( p 0 , π θ , p ) [ b ( s t ) ∇ θ log p θ ( a t ∣ s t ) ] = E s 0 ∼ p 0 [ b ( s t ) E a t ∼ π θ ( ⋅ ∣ s t ) [ ∇ θ log π θ ( a t ∣ s t ) ] ] = E s 0 ∼ p 0 [ b ( s t ) a ∑ ∇ θ π θ ( a ∣ s ) ] = E s 0 ∼ p 0 [ b ( s t ) ∇ θ a ∑ π θ ( a ∣ s ) ] = E s 0 ∼ p 0 [ b ( s t ) ∇ θ 1 ] = 0 因此,policy gradient 满足 baseline invariance property.
另一方面,Rao–Blackwell Theorem 说明了我们可以通过构造一个新的无偏估计,这个新的无偏估计相对于原始的无偏估计,有更小的方差。
Theorem: Rao-Blackwell Theorem 令 X X X 和 Y Y Y 为随机变量, 令 I ^ 1 ( X , Y ) \hat{I}_1(X,Y) I ^ 1 ( X , Y ) 为 I I I 的一个无偏估计,即
I = E X Y [ I ^ 1 ( X Y ) ] I = \mathbb{E}_{XY}[\hat{I}_1(XY)] I = E X Y [ I ^ 1 ( X Y )] 令 I ^ 2 ( Y ) = E X ∣ Y [ I ^ 1 ( X , Y ) ∣ Y ] \hat{I}_2(Y)=\mathbb{E}_{X\mid Y}[\hat{I}_1(X,Y)\mid Y] I ^ 2 ( Y ) = E X ∣ Y [ I ^ 1 ( X , Y ) ∣ Y ] , 则 I ^ 2 \hat{I}_2 I ^ 2 也是 I I I 的一个无偏估计,并且
v a r [ I ^ 2 ] ≤ v a r [ I ^ 1 ] \mathrm{var}[\hat{I}_2]\leq\mathrm{var}[\hat{I}_1] var [ I ^ 2 ] ≤ var [ I ^ 1 ] I ^ 2 \hat{I}_2 I ^ 2 被称为 I ^ 1 ( X , Y ) \hat{I}_1(X,Y) I ^ 1 ( X , Y ) 的Rao-Blackwellized estimator .
Proof: 由 Expectation 中的 law of total expectation, 我们有
E X Y [ I ^ 2 ( X Y ) ] = E Y [ E X ∣ Y [ I ^ 1 ( X , Y ) ∣ Y ] ] = E X Y [ I ^ 1 ( X Y ) ] = 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 E X Y [ I ^ 2 ( X Y )] = E Y [ E X ∣ Y [ I ^ 1 ( X , Y ) ∣ Y ]] = E X Y [ I ^ 1 ( X Y )] = I 其次,我们有
v a r [ I ^ 2 ] = E Y [ ( I − I ^ 2 ( Y ) ) 2 ] = E Y [ ( E X ∣ Y [ I − I ^ 1 ( X , Y ) ∣ Y ] ) 2 ] ≤ E Y [ ( E X ∣ Y [ ( I − I ^ 1 ( X , Y ) ) 2 ∣ Y ] ) ] = E Y [ ( I − I ^ 1 ( X , Y ) ) 2 ] = v a r [ 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} var [ I ^ 2 ] = E Y [( I − I ^ 2 ( Y ) ) 2 ] = E Y [ ( E X ∣ Y [ I − I ^ 1 ( X , Y ) ∣ Y ] ) 2 ] ≤ E Y [ ( E X ∣ Y [ ( I − I ^ 1 ( X , Y ) ) 2 ∣ Y ] ) ] = E Y [( I − I ^ 1 ( X , Y ) ) 2 ] = var [ I ^ 1 ] 这里不等式使用了 Jensen’s inequality. 证毕。■ \blacksquare ■
基于 Rao–Blackwell Theorem, 我们可以通过改变 b ( s t ) b(s_t) b ( s t ) 来调整 policy gradient 的 variance.
Reward to go 的基本思想为
对于时刻 t t t , 过去的轨迹 s 0 , r 0 , … , r t − 1 s_0,r_0,\dots,r_{t - 1} s 0 , r 0 , … , r t − 1
不应该被考虑,因为它们都已经固定了,加入它们会产生额外的 variance,
此时,我们对应的 b ( s t ) b(s_t) b ( s t ) 定义为
其定义的 baseline 如下所示
b ( s t ) = ∑ k = 0 t γ k r k b(s_t) = \sum_{k=0}^{t}\gamma^k r_k b ( s t ) = k = 0 ∑ t γ k r k
此时,我们的 policy gradient 为
∇ θ J ( θ ) = E τ ∼ ( p 0 , π θ , p ) [ ∑ t = 0 T − 1 ∇ θ log π θ ( a t ∣ s t ) ∑ k = t T − 1 r ( s k ) ] \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] ∇ θ J ( θ ) = E τ ∼ ( p 0 , π θ , p ) [ t = 0 ∑ T − 1 ∇ θ log π θ ( a t ∣ s t ) k = t ∑ T − 1 r ( s k ) ]
对应的 REINFORCE 算法改进版本为
Algorithm: Policy Gradient with MC and reward to go (REINFORCE)
Initialize policy parameters θ 0 \theta_0 θ 0 .
for k = 0 , 1 , … k=0,1,\dots k = 0 , 1 , … :
Sample M M M trajectories { τ i } \{\tau_i\} { τ i } according to policy π θ k \pi_{\theta_k} π θ k .
Estimate the gradient via MC:
g ^ = 1 M ∑ i = 1 M ∇ θ log π θ k ( a t ∣ s t ) γ t G t ( τ 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) g ^ = M 1 i = 1 ∑ M ∇ θ log π θ k ( a t ∣ s t ) γ t G t ( τ i )
Update θ k \theta_k θ k using g ^ \hat{g} g ^ with an optimizer.
我们可以从理论上推导出最优 baseline.
Lemma: 令 s s s 和 a a a 为随机变量,w ( s , a ) w(s,a) w ( s , a ) , Q ( s , a ) Q(s,a) Q ( s , a ) , b ( s ) b(s) b ( s ) 为函数,则
b ∗ = arg min b ( ⋅ ) E a , s [ w 2 ( 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 ∗ = arg b ( ⋅ ) min E a , s [ w 2 ( s , a ) ( Q ( s , a ) − b ( s ) ) 2 ] 其中
b ∗ ( s ) = E a ∣ s [ w 2 ( s , a ) Q ( s , a ) ∣ s ] E a ∣ s [ w 2 ( 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]} b ∗ ( s ) = E a ∣ s [ w 2 ( s , a ) ∣ s ] E a ∣ s [ w 2 ( s , a ) Q ( s , a ) ∣ s ]
Proof: 对目标函数进行展开得到
E a , s [ w 2 ( s , a ) ( Q ( s , a ) − b ( s ) ) 2 ] = E a , s [ w 2 ( s , a ) ( Q ( s , a ) − b ∗ ( s ) + b ∗ ( s ) − b ( s ) ) 2 ] = E a , s [ w 2 ( s , a ) ( Q ( s , a ) − b ∗ ( s ) ) 2 ] + E a , s [ w 2 ( s , a ) ( b ∗ ( s ) − b ( s ) ) 2 ] + 2 E a , s [ w 2 ( s , a ) ( Q ( s , a ) − b ∗ ( s ) ) ( b ∗ ( s ) − b ( s ) ) ] = E a , s [ w 2 ( s , a ) ( Q ( s , a ) − b ∗ ( s ) ) 2 ] + E a , s [ w 2 ( s , a ) ( b ∗ ( s ) − b ( s ) ) 2 ] ≥ E a , s [ w 2 ( 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} E a , s [ w 2 ( s , a ) ( Q ( s , a ) − b ( s ) ) 2 ] = E a , s [ w 2 ( s , a ) ( Q ( s , a ) − b ∗ ( s ) + b ∗ ( s ) − b ( s ) ) 2 ] = E a , s [ w 2 ( s , a ) ( Q ( s , a ) − b ∗ ( s ) ) 2 ] + E a , s [ w 2 ( s , a ) ( b ∗ ( s ) − b ( s ) ) 2 ] + 2 E a , s [ w 2 ( s , a ) ( Q ( s , a ) − b ∗ ( s ) ) ( b ∗ ( s ) − b ( s ) ) ] = E a , s [ w 2 ( s , a ) ( Q ( s , a ) − b ∗ ( s ) ) 2 ] + E a , s [ w 2 ( s , a ) ( b ∗ ( s ) − b ( s ) ) 2 ] ≥ E a , s [ w 2 ( s , a ) ( Q ( s , a ) − b ∗ ( s ) ) 2 ] 其中第三个等式用到了
E a , s [ w 2 ( s , a ) ( Q ( s , a ) − b ∗ ( s ) ) ( b ∗ ( s ) − b ( s ) ) ] = E s [ E a ∣ s [ w 2 ( s , a ) ( Q ( s , a ) − b ∗ ( s ) ) ( b ∗ ( s ) − b ( s ) ) ] ∣ s ] = E s [ ( E a ∣ s [ w 2 ( s , a ) Q ( s , a ) ∣ s ] − E a ∣ s [ w 2 ( 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} E a , s [ w 2 ( s , a ) ( Q ( s , a ) − b ∗ ( s ) ) ( b ∗ ( s ) − b ( s ) ) ] = E s [ E a ∣ s [ w 2 ( s , a ) ( Q ( s , a ) − b ∗ ( s ) ) ( b ∗ ( s ) − b ( s ) ) ] ∣ s ] = E s [ ( E a ∣ s [ w 2 ( s , a ) Q ( s , a ) ∣ s ] − E a ∣ s [ w 2 ( s , a ) ∣ s ] b ∗ ( s ) ) ( b ∗ ( s ) − b ( s ) ) ] = 0 证毕。■ \blacksquare ■
基于上述引理,我们可以得到关于 baseline invariance 中的最优 baseline:
b ∗ ( s ) = E a t ∼ π θ ( ⋅ ∣ s t ) [ ( ∇ θ log π θ ( a t ∣ s t ) ) 2 Q π θ ( s , a ) ∣ s ] E a t ∼ π θ ( ⋅ ∣ s t ) [ ∇ θ log π θ ( a t ∣ s t ) ) 2 ∣ s ] 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]} b ∗ ( s ) = E a t ∼ π θ ( ⋅ ∣ s t ) [ ∇ θ log π θ ( a t ∣ s t ) ) 2 ∣ s ] E a t ∼ π θ ( ⋅ ∣ s t ) [ ( ∇ θ log π θ ( a t ∣ s t ) ) 2 Q π θ ( s , a ) ∣ s ]
理论上我们可以使用 b ∗ ( s t ) b^*(s_t) b ∗ ( s t ) 作为 baseline,但是实际上由于其表达式太复杂,我们一般使用其简化版本。
简化版本移除了 b ∗ ( s ) b^*(s) b ∗ ( s ) 中的 ∇ θ log π θ ( a t ∣ s t ) \nabla_\theta\log\pi_\theta(a_t\mid s_t) ∇ θ log π θ ( a t ∣ s t ) , 这样我们的估计就是
b ( s ) = E a t ∼ π θ ( ⋅ ∣ s t ) [ Q π θ ( s , a ) ∣ s ] E a t ∼ π θ ( ⋅ ∣ s t ) [ 1 ∣ s ] = 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 ) = E a t ∼ π θ ( ⋅ ∣ s t ) [ 1 ∣ s ] E a t ∼ π θ ( ⋅ ∣ s t ) [ Q π θ ( s , a ) ∣ s ] = V π θ ( s )
b ( s ) = V π θ ( s ) b(s)=V^{\pi_\theta}(s) b ( s ) = V π θ ( s ) 并不是最优 baseline, 但是其性质比较好。
我们将 b ( s ) = V π θ ( s ) b(s)=V^{\pi_\theta}(s) b ( s ) = V π θ ( s ) 带入到目标函数中就得到经典的 policy gradient 表达式
∇ θ J ( θ ) = E τ ∼ ( p 0 , π θ , p ) [ ∑ t = 0 T − 1 ∇ θ log π θ ( a t ∣ s t ) ( Q π θ ( s t , a t ) − V π θ ( s t ) ) ] \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] ∇ θ J ( θ ) = E τ ∼ ( p 0 , π θ , p ) [ t = 0 ∑ T − 1 ∇ θ log π θ ( a t ∣ s t ) ( Q π θ ( s t , a t ) − V π θ ( s t ) ) ]
我们定义
A π θ ( s t , a t ) : = Q π θ ( s t , a t ) − V π θ ( s t ) \boxed{
A^{\pi_\theta}(s_t,a_t) := Q^{\pi_\theta}(s_t,a_t)-V^{\pi_\theta}(s_t)
} A π θ ( s t , a t ) := Q π θ ( s t , a t ) − V π θ ( s t )
为 a t a_t a t 在 s t s_t s t 处的 advantage . 注意到 V π θ ( s t ) V^{\pi_\theta}(s_t) V π θ ( s t ) 评估了当前状态下沿着策略 π θ \pi_\theta π θ 所能够获得的 discounted return, 因此 advantage 说明了 action a t a_t a t 对于 π θ \pi_\theta π θ 的平均好坏程度:
如果 A π θ ( s t , a t ) > 0 A^{\pi_\theta}(s_t,a_t) >0 A π θ ( s t , a t ) > 0 , 则说明 a t a_t a t (平均意义下)相对于 π θ \pi_\theta π θ 来说是一个更好的 action
反之如果 A π θ ( s t , a t ) < 0 A^{\pi_\theta}(s_t,a_t) <0 A π θ ( s t , a t ) < 0 , 则说明 a t a_t a t (平均意义下)相对于 π θ \pi_\theta π θ 来说是一个更坏的 action
实际场景下,由于 V π θ ( s t ) V^{\pi_\theta}(s_t) V π θ ( s t ) 未知,使用前面的 value iteration 方法,我们可以使用估计
A ^ π θ ( s t , a t ) : = Q π θ ( s t , a t ) − V ϕ ( s t ) \hat{A}^{\pi_\theta}(s_t,a_t) := Q^{\pi_\theta}(s_t,a_t)-V^{\phi}(s_t) A ^ π θ ( s t , a t ) := Q π θ ( s t , a t ) − V ϕ ( s t )
来代替。
对于最优策略 π ∗ \pi^* π ∗ ,我们有
A π ∗ ( s t , a t ) : = Q π ∗ ( s t , a t ) − V π ∗ ( s t ) ≤ 0 A^{\pi^*}(s_t,a_t) := Q^{\pi^*}(s_t,a_t)-V^{\pi^*}(s_t)\leq 0 A π ∗ ( s t , a t ) := Q π ∗ ( s t , a t ) − V π ∗ ( s t ) ≤ 0
证明也很简单,注意到 V π ∗ ( s t ) = max a Q π ∗ ( s t , a t ) V^{\pi^*}(s_t)=\max_{a}Q^{\pi^*}(s_t,a_t) V π ∗ ( s t ) = max a Q π ∗ ( s t , a t ) 即可。
最后,我们的 REINFORCE 算法改进如下
Algorithm: Policy Gradient with MC and value function (REINFORCE)
Initialize policy parameters θ 0 \theta_0 θ 0 .
for k = 0 , 1 , … k=0,1,\dots k = 0 , 1 , … :
Sample M M M trajectories { τ i } \{\tau_i\} { τ i } according to policy π θ k \pi_{\theta_k} π θ k .
Estimate the gradient via MC:
g ^ = 1 M ∑ i = 1 M ∇ θ log π θ k ( a t ∣ s t ) γ t ( Q π θ ( s t , a t ) − V ϕ ( s t ) ) \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)) g ^ = M 1 i = 1 ∑ M ∇ θ log π θ k ( a t ∣ s t ) γ t ( Q π θ ( s t , a t ) − V ϕ ( s t ))
Update θ k \theta_k θ k using g ^ \hat{g} g ^ with an optimizer.
最后,我们把前面的 basline 使用统一的公式进行表示,记
∇ θ J ( θ ) = E τ ∼ ( p 0 , π θ , p ) [ Ψ ( τ ) ∇ θ log p θ ( τ ) ] = E τ ∼ ( p 0 , π θ , p ) [ ∑ t = 0 T Ψ t ( τ ) ∇ θ log p θ ( τ ) ] \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] ∇ θ J ( θ ) = E τ ∼ ( p 0 , π θ , p ) [ Ψ ( τ ) ∇ θ log p θ ( τ ) ] = E τ ∼ ( p 0 , π θ , p ) [ t = 0 ∑ T Ψ t ( τ ) ∇ θ log p θ ( τ ) ]
则不同的方法对应的表达式如下
TODO
上一节使用完整 trajectory return 构造 REINFORCE gradient。这个 estimator
不需要学习 value function,但在长序列和稀疏 reward 下通常具有较高 variance。
baseline invariance 允许我们用 advantage 替代原始 return,而不改变期望梯度。
本节把前面的 policy evaluation 与 policy gradient 合并:actor
π θ \pi_\theta π θ 负责生成和更新策略,critic V ϕ V_\phi V ϕ 或 Q ϕ Q_\phi Q ϕ 负责估计
return/advantage。对 LLM 而言,两者分别对应需要优化的语言模型和训练期间使用的
value estimator。
我们首先考虑最简单的 actor-critic algorithm, 注意到
∇ θ J ( θ ) = E τ ∼ ( p 0 , π θ , p ) [ ∑ t = 0 T − 1 ∇ θ log π θ ( a t ∣ s t ) Q π θ ( s t , a t ) ] \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] ∇ θ J ( θ ) = E τ ∼ ( p 0 , π θ , p ) [ t = 0 ∑ T − 1 ∇ θ log π θ ( a t ∣ s t ) Q π θ ( s t , a t ) ]
因此,其对应的 policy gradient methods 为
θ k + 1 = θ k + α E τ ∼ ( p 0 , π θ , p ) [ ∑ t = 0 T − 1 ∇ θ log π θ ( a t ∣ s t ) Q π θ ( s t , a t ) ] \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] θ k + 1 = θ k + α E τ ∼ ( p 0 , π θ , p ) [ t = 0 ∑ T − 1 ∇ θ log π θ ( a t ∣ s t ) Q π θ ( s t , a t ) ]
在前面的章节中,我们介绍了基于 MC 和 TD 两种方式来估计
Q π θ ( s t , a t ) Q^{\pi_\theta}(s_t,a_t) Q π θ ( s t , a t ) 。据此可以区分:
如果 Q π θ ( s t , a t ) Q^{\pi_\theta}(s_t,a_t) Q π θ ( s t , a t ) 由 MC 来进行估计,则我们将其称为 REINFORCE 或者 Monte Carlo policy gradient.
如果 Q π θ ( s t , a t ) Q^{\pi_\theta}(s_t,a_t) Q π θ ( s t , a t ) 由 TD learning 来进行估计,则我们将其称为 actor-critic , 这是我们本节的重点介绍内容
最简单的 actor-critic algorithm 如下所示
Algorithm: Q Actor-Critic
Initialize policy parameters θ 0 \theta_0 θ 0 , value function parameters ϕ 0 \phi_0 ϕ 0 .
for t = 0 , 1 , … , T − 1 t = 0, 1, \dots, T-1 t = 0 , 1 , … , T − 1 :
a t ∼ π θ ( ⋅ ∣ s t ) a_t \sim \pi_{\theta}(\cdot \mid s_t) a t ∼ π θ ( ⋅ ∣ s t ) , r t , s t + 1 ∼ p ( ⋅ , ⋅ ∣ s t , a t ) r_t, s_{t+1} \sim p(\cdot, \cdot \mid s_t, a_t) r t , s t + 1 ∼ p ( ⋅ , ⋅ ∣ s t , a t ) , a t + 1 ∼ π θ ( ⋅ ∣ s t + 1 ) a_{t+1} \sim \pi_{\theta}(\cdot \mid s_{t+1}) a t + 1 ∼ π θ ( ⋅ ∣ s t + 1 ) .
(ACTOR) policy update :
θ t + 1 = θ t + α θ ∇ θ log π θ ( a t ∣ s t ) Q ϕ t ( s t , a t ) \theta_{t+1} = \theta_t + \alpha_{\theta} \nabla_{\theta} \log \pi_{\theta}(a_t \mid s_t) \, Q^{\phi_t}(s_t, a_t) θ t + 1 = θ t + α θ ∇ θ log π θ ( a t ∣ s t ) Q ϕ t ( s t , a t )
(CRITIC) value update :
ϕ t + 1 = ϕ t + α ϕ [ r t + γ Q ϕ t ( s t + 1 , a t + 1 ) − Q ϕ t ( s t , a t ) ] ∇ ϕ Q ϕ t ( s t , a t ) \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) ϕ t + 1 = ϕ t + α ϕ [ r t + γ Q ϕ t ( s t + 1 , a t + 1 ) − Q ϕ t ( s t , a t ) ] ∇ ϕ Q ϕ t ( s t , a t )
在上一节中,我们介绍了 QAC, 但实际上我们用的更多的是 advantage actor-critic, 其梯度如下所示
∇ θ J ( θ ) = E τ ∼ ( p 0 , π θ , p ) [ ∑ t = 0 T − 1 ∇ θ log π θ ( a t ∣ s t ) A π θ ( s t , a t ) ] \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] ∇ θ J ( θ ) = E τ ∼ ( p 0 , π θ , p ) [ t = 0 ∑ T − 1 ∇ θ log π θ ( a t ∣ s t ) A π θ ( s t , a t ) ]
注意到
A π θ ( s t , a t ) = Q π θ ( s t , a t ) − V π θ ( s t ) = E π θ [ r t + γ V π θ ( s t + 1 ) − V π θ ( s t ) ∣ s t , a t ] 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] A π θ ( s t , a t ) = Q π θ ( s t , a t ) − V π θ ( s t ) = E π θ [ r t + γ V π θ ( s t + 1 ) − V π θ ( s t ) ∣ s t , a t ]
因此我们可以用 TD error 来进行近似:
A ^ t = r t + γ V ϕ t ( s t + 1 ) − V ϕ t ( s t ) \hat{A}_t= r_t + \gamma V_{\phi_t}(s_{t+1})-V_{\phi_t}(s_t) A ^ t = r t + γ V ϕ t ( s t + 1 ) − V ϕ t ( s t )
这样我们就得到了 A2C 算法
Algorithm: Advantage Actor-Critic (A2C)
Initialize policy parameters θ 0 \theta_0 θ 0 , value function parameters ϕ 0 \phi_0 ϕ 0 .
for t = 0 , 1 , … , T − 1 t = 0, 1, \dots, T-1 t = 0 , 1 , … , T − 1 :
a t ∼ π θ ( ⋅ ∣ s t ) a_t \sim \pi_{\theta}(\cdot \mid s_t) a t ∼ π θ ( ⋅ ∣ s t ) , r t , s t + 1 ∼ p ( ⋅ , ⋅ ∣ s t , a t ) r_t, s_{t+1} \sim p(\cdot, \cdot \mid s_t, a_t) r t , s t + 1 ∼ p ( ⋅ , ⋅ ∣ s t , a t ) .
Advantage estimation :
A ^ t ≈ r t + γ V ϕ t ( s t + 1 ) − V ϕ t ( s t ) \hat{A}_t \approx r_t + \gamma V_{\phi_t}(s_{t+1}) - V_{\phi_t}(s_t) A ^ t ≈ r t + γ V ϕ t ( s t + 1 ) − V ϕ t ( s t )
(ACTOR) policy update :
θ t + 1 = θ t + α θ ∇ θ log π θ ( a t ∣ s t ) A ^ t \theta_{t+1} = \theta_t + \alpha_{\theta} \nabla_{\theta} \log \pi_{\theta}(a_t \mid s_t) \, \hat{A}_t θ t + 1 = θ t + α θ ∇ θ log π θ ( a t ∣ s t ) A ^ t
(CRITIC) value update :
ϕ t + 1 = ϕ t + α ϕ A ^ t ∇ ϕ V ϕ t ( s t ) \phi_{t+1} = \phi_t + \alpha_{\phi} \, \hat{A}_t \, \nabla_{\phi} V^{\phi_t}(s_t) ϕ t + 1 = ϕ t + α ϕ A ^ t ∇ ϕ V ϕ t ( s t )
在上一节中,我们介绍了我们可以用 1-step TD error 来近似 advantage,
对比结果如下表所示
为了实现更好的 bias-variance trade-off, GAE (Schulman et al., 2016 ) 提出了使用一个参数 λ ∈ [ 0 , 1 ] \lambda\in[0,1] λ ∈ [ 0 , 1 ] 来控制 bias 和 variance 之间的权重,其表达式如下所示
A ^ t G A E ( γ , λ ) = ∑ ℓ = 0 ∞ ( γ λ ) ℓ δ t + ℓ \hat{A}_t^{GAE(\gamma,\lambda)} = \sum_{\ell=0}^{\infty}(\gamma\lambda)^{\ell}\delta_{t+\ell} A ^ t G A E ( γ , λ ) = ℓ = 0 ∑ ∞ ( γ λ ) ℓ δ t + ℓ
其中
δ t = r t + γ V ϕ t ( s t + 1 ) − V ϕ t ( s t ) \delta_t = r_t + \gamma V_{\phi_t}(s_{t+1})-V_{\phi_t}(s_t) δ t = r t + γ V ϕ t ( s t + 1 ) − V ϕ t ( s t )
我们可以探究不同的 λ \lambda λ 对应的估计:
当 λ = 0 \lambda=0 λ = 0 时,我们有
A ^ t G A E ( γ , 0 ) = ( γ ⋅ 0 ) 0 δ t = δ t \hat{A}_t^{GAE(\gamma,0)} = (\gamma\cdot0)^0\delta_t=\delta_t A ^ t G A E ( γ , 0 ) = ( γ ⋅ 0 ) 0 δ t = δ t
此时,GAE 就退化成了 one-step TD error estimate,
当 λ = 1 \lambda=1 λ = 1 时,我们有
A ^ t G A E ( γ , 1 ) = ∑ ℓ = 0 ∞ γ ℓ δ t + ℓ = ∑ ℓ = 0 ∞ γ ℓ ( r t + γ V ϕ t ( s t + 1 ) − V ϕ t ( s t ) ) \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) A ^ t G A E ( γ , 1 ) = ℓ = 0 ∑ ∞ γ ℓ δ t + ℓ = ℓ = 0 ∑ ∞ γ ℓ ( r t + γ V ϕ t ( s t + 1 ) − V ϕ t ( s t ) )
如果我们假设这里使用的是真实的 value function 的话,那么我们有
A ^ t G A E ( γ , 1 ) ≈ ∑ ℓ = 0 ∞ γ ℓ r t + ℓ + 1 − V ( s t ) = Q π θ ( s t , a t ) − V ( s t ) \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) A ^ t G A E ( γ , 1 ) ≈ ℓ = 0 ∑ ∞ γ ℓ r t + ℓ + 1 − V ( s t ) = Q π θ ( s t , a t ) − V ( s t )
当 0 < λ < 1 0<\lambda<1 0 < λ < 1 时,G A E GAE G A E 是 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)
Given policy parameters θ \theta θ , value function parameters ϕ \phi ϕ .
Sample a trajectory τ ∼ ( p 0 , π θ , p ) \tau \sim (p_0, \pi_{\theta}, p) τ ∼ ( p 0 , π θ , p ) .
for t = 0 , 1 , … , T − 1 t = 0, 1, \dots, T-1 t = 0 , 1 , … , T − 1 :
Compute TD errors:
δ t = r t + γ V ϕ ( s t + 1 ) − V ϕ ( s t ) \delta_t = r_t + \gamma V_{\phi}(s_{t+1}) - V_{\phi}(s_t) δ t = r t + γ V ϕ ( s t + 1 ) − V ϕ ( s t )
Initialize A ^ T G A E = 0 \hat{A}_{T}^{\mathrm{GAE}} = 0 A ^ T GAE = 0 .
for t = T − 1 , … , 0 t = T-1, \dots, 0 t = T − 1 , … , 0 :
Iterate backward:
A ^ t G A E = δ t + γ λ A ^ t + 1 G A E \hat{A}_t^{\mathrm{GAE}} = \delta_t + \gamma\lambda \, \hat{A}_{t+1}^{\mathrm{GAE}} A ^ t GAE = δ t + γ λ A ^ t + 1 GAE
GAE 的主要优点在于其相对于 MC estimate 可以大幅度降低 variance, 从而提高训练的稳定性以及效率
comparison with value-based methods:
无法处理高维或者连续动作空间,actor-critic 中 actor 可以输出一个连续值的概率分布,直接从分布中采样就能得到动作,不需要遍历动作空间
无法学习到随机策略: value-based methods 只能推导出确定性策略,使用 ϵ \epsilon ϵ greedy 之后,也只是随机探索,而不是有策略的随机探索。Actor 可以参数化策略,学习到任意概率分布的随机策略
学习到的策略更平滑:value-based methods 动作选择依赖于 greedy strategy, 导致 Q function 微小变化导致策略发生剧烈跳变。而 actor-critic 通过 policy gradient 解决了这个问题
comparison with policy-based methods
policy-based methods 依赖完整的 rollout 的 return 来更新,对于 long-term tasks, 最后的 return 方差很大,训练难以收敛。 actor-critic 通过引入 critic 来估计期望收益,计算 advantage, 来降低整体法国差,同时用 TD error 代替回合更新,加速训练。
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) .
前面的 policy gradient 和 actor-critic 给出了更新方向,但没有回答一个关键问题:
一次可以把 policy 更新多远? 如果新旧 policy 差异过大,旧 policy 生成的
trajectory 很快失效,近似目标也可能不再可靠。
TRPO (Trust Region Policy Optimization) (Schulman et al., 2015 ) 使用旧策略
π θ o l d \pi_{\theta_{\mathrm{old}}} π θ 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 ( θ ) = E s 0 ∼ p 0 [ V π θ ( s 0 ) ] \max_{\pi}\mathcal{J}(\theta)=\mathbb{E}_{s_0\sim p_0}\left[V^{\pi_\theta}(s_0)\right] π max J ( θ ) = E s 0 ∼ p 0 [ V π θ ( s 0 ) ]
我们假设有两个 policy 的参数, θ \theta θ 和 θ o l d \theta_{old} θ o l d , 我们分析一下两个 policy 对应的目标函数之间的 difference:
J ( θ ) − J ( θ o l d ) = J ( θ ) − E s 0 ∼ p 0 [ V π θ o l d ( s 0 ) ] = J ( θ ) − E τ ∼ ( p 0 , π θ , p ) [ V π θ o l d ( s 0 ) ] = J ( θ ) − E τ ∼ ( p 0 , π θ , p ) [ ∑ t = 0 T − 1 γ t V π θ o l d ( s t ) − ∑ t = 1 T − 1 γ t V π θ o l d ( s t ) ] = E τ ∼ ( p 0 , π θ , p ) [ ∑ t = 0 T − 1 γ t r t ] + E τ ∼ ( p 0 , π θ , p ) [ ∑ t = 0 T − 1 γ t ( γ V π θ o l d ( s t + 1 ) − V π θ o l d ( s t ) ) ] = E τ ∼ ( p 0 , π θ , p ) [ ∑ t = 0 T − 1 γ t ( r t + γ V π θ o l d ( s t + 1 ) − V π θ o l d ( s t ) ) ] = E τ ∼ ( p 0 , π θ , p ) [ ∑ t = 0 T − 1 γ t ( γ Q π θ o l d ( s t , a t ) − V π θ o l d ( s t ) ) ] = E τ ∼ ( p 0 , π θ , p ) [ ∑ t = 0 T − 1 γ t A π θ o l d ( s t , a t ) ] \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 ( θ o l d ) = J ( θ ) − E s 0 ∼ p 0 [ V π θ o l d ( s 0 ) ] = J ( θ ) − E τ ∼ ( p 0 , π θ , p ) [ V π θ o l d ( s 0 ) ] = J ( θ ) − E τ ∼ ( p 0 , π θ , p ) [ t = 0 ∑ T − 1 γ t V π θ o l d ( s t ) − t = 1 ∑ T − 1 γ t V π θ o l d ( s t ) ] = E τ ∼ ( p 0 , π θ , p ) [ t = 0 ∑ T − 1 γ t r t ] + E τ ∼ ( p 0 , π θ , p ) [ t = 0 ∑ T − 1 γ t ( γ V π θ o l d ( s t + 1 ) − V π θ o l d ( s t ) ) ] = E τ ∼ ( p 0 , π θ , p ) [ t = 0 ∑ T − 1 γ t ( r t + γ V π θ o l d ( s t + 1 ) − V π θ o l d ( s t ) ) ] = E τ ∼ ( p 0 , π θ , p ) [ t = 0 ∑ T − 1 γ t ( γ Q π θ o l d ( s t , a t ) − V π θ o l d ( s t ) ) ] = E τ ∼ ( p 0 , π θ , p ) [ t = 0 ∑ T − 1 γ t A π θ o l d ( s t , a t ) ]
我们进一步展开得到
J ( θ ) − J ( θ o l d ) = E τ ∼ ( p 0 , π θ , p ) [ ∑ t = 0 T − 1 γ t A π θ o l d ( s t , a t ) ] = E s 0 ∼ p 0 , a 0 ∼ π θ ( ⋅ ∣ s 0 ) [ A π θ o l d ( s 0 , a 0 ) + E π θ [ ∑ t = 1 T − 1 γ t A π θ o l d ( s t , a t ) ∣ s 0 , a 0 ] ] = E s 0 ∼ p 0 , a 0 ∼ π θ ( ⋅ ∣ s 0 ) , a 0 ′ ∼ π θ o l d ( ⋅ ∣ s 0 ) [ π θ ( a 0 ′ ∣ s 0 ) π θ o l d ( a 0 ′ ∣ s 0 ) A π θ o l d ( s 0 , a 0 ′ ) + E π θ [ ∑ t = 1 T − 1 γ t A π θ o l d ( s t , a t ) ∣ s 0 , a 0 ] ] = … ( keep enrolling the summation ) = E τ ∼ ( p 0 , π θ , p ) , a t ′ ∼ π θ o l d ( ⋅ ∣ s t ) [ ∑ t = 0 T − 1 γ t π θ ( a t ′ ∣ s t ) π θ o l d ( a t ′ ∣ s t ) A π θ o l d ( s t , a t ′ ) ] \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 ( θ ) − J ( θ o l d ) = E τ ∼ ( p 0 , π θ , p ) [ t = 0 ∑ T − 1 γ t A π θ o l d ( s t , a t ) ] = E s 0 ∼ p 0 , a 0 ∼ π θ ( ⋅ ∣ s 0 ) [ A π θ o l d ( s 0 , a 0 ) + E π θ [ t = 1 ∑ T − 1 γ t A π θ o l d ( s t , a t ) ∣ s 0 , a 0 ] ] = E s 0 ∼ p 0 , a 0 ∼ π θ ( ⋅ ∣ s 0 ) , a 0 ′ ∼ π θ o l d ( ⋅ ∣ s 0 ) [ π θ o l d ( a 0 ′ ∣ s 0 ) π θ ( a 0 ′ ∣ s 0 ) A π θ o l d ( s 0 , a 0 ′ ) + E π θ [ t = 1 ∑ T − 1 γ t A π θ o l d ( s t , a t ) ∣ s 0 , a 0 ] ] = … ( keep enrolling the summation ) = E τ ∼ ( p 0 , π θ , p ) , a t ′ ∼ π θ o l d ( ⋅ ∣ s t ) [ t = 0 ∑ T − 1 γ t π θ o l d ( a t ′ ∣ s t ) π θ ( a t ′ ∣ s t ) A π θ o l d ( s t , a t ′ ) ]
因此,我们有
J ( θ ) = E τ ∼ ( p 0 , π θ , p ) , a t ′ ∼ π θ o l d ( ⋅ ∣ s t ) [ ∑ t = 0 T − 1 γ t π θ ( a t ′ ∣ s t ) π θ o l d ( a t ′ ∣ s t ) A π θ o l d ( s t , a t ′ ) ] + 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 J ( θ ) = E τ ∼ ( p 0 , π θ , p ) , a t ′ ∼ π θ o l d ( ⋅ ∣ s t ) [ t = 0 ∑ T − 1 γ t π θ o l d ( a t ′ ∣ s t ) π θ ( a t ′ ∣ s t ) A π θ o l d ( s t , a t ′ ) ] + C
但是,现在我们的更姓还是依赖于 π θ \pi_\theta π θ , 这很难处理。因此 TRPO 的思想就是,当 π θ \pi_\theta π θ 与 π θ o l d \pi_{\theta_{old}} π θ o l d 充分接近时,使用 π θ o l d \pi_{\theta_{old}} π θ o l d 来替换 π θ \pi_\theta π θ , 这样,上面目标函数就变成了
max θ K ( θ ; θ o l d ) = E τ ∼ ( p 0 , π θ o l d , p ) , a t ′ ∼ π θ o l d ( ⋅ ∣ s t ) [ ∑ t = 0 T − 1 γ t π θ ( a t ′ ∣ s t ) π θ o l d ( a t ′ ∣ s t ) A π θ o l d ( s t , a t ′ ) ] s . t . π θ o l d ≈ π θ \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} θ max s.t. K ( θ ; θ o l d ) = E τ ∼ ( p 0 , π θ o l d , p ) , a t ′ ∼ π θ o l d ( ⋅ ∣ s t ) [ t = 0 ∑ T − 1 γ t π θ o l d ( a t ′ ∣ s t ) π θ ( a t ′ ∣ s t ) A π θ o l d ( s t , a t ′ ) ] π θ o l d ≈ π θ
我们可以证明
∇ J ( θ ) ∣ θ = θ o l d = ∇ θ K ( θ ; θ o l d ) ∣ θ = θ o l d \nabla \mathcal{J}(\theta)\mid _{\theta=\theta_{old}} = \nabla_\theta \mathcal{K}(\theta;\theta_{old})\mid _{\theta=\theta_{old}} ∇ J ( θ ) ∣ θ = θ o l d = ∇ θ K ( θ ; θ o l d ) ∣ θ = θ o l d
这就是 TRPO 的核心改进,在实际求解时,我们将上面的问题规范为如下形式
max θ K ( θ ; θ o l d ) = E τ ∼ ( p 0 , π θ o l d , p ) , a t ′ ∼ π θ o l d ( ⋅ ∣ s t ) [ ∑ t = 0 T − 1 γ t π θ ( a t ′ ∣ s t ) π θ o l d ( a t ′ ∣ s t ) A π θ o l d ( s t , a t ′ ) ] s . t . max s ∈ S K L ( π θ o l d ( ⋅ ∣ 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} θ max s.t. K ( θ ; θ o l d ) = E τ ∼ ( p 0 , π θ o l d , p ) , a t ′ ∼ π θ o l d ( ⋅ ∣ s t ) [ t = 0 ∑ T − 1 γ t π θ o l d ( a t ′ ∣ s t ) π θ ( a t ′ ∣ s t ) A π θ o l d ( s t , a t ′ ) ] s ∈ S max KL ( π θ o l d ( ⋅ ∣ s ) ∣∣ π θ ( ⋅ ∣ s )) ≤ δ
这里 δ > 0 \delta>0 δ > 0 为一个超参数。
现在我们可以采样多条轨迹 τ ( i ) ∼ ( p 0 , π θ o l d , p ) \tau^{(i)}\sim (p_0,\pi_{\theta_{old}}, p) τ ( i ) ∼ ( p 0 , π θ o l d , p ) 再使用 MC 方法来估计目标函数 K ( θ ; θ o l d ) \mathcal{K}(\theta;\theta_{old}) K ( θ ; θ o l d )
K ( θ ; θ o l d ) = E τ ∼ ( p 0 , π θ o l d , p ) , a t ∼ π θ o l d ( ⋅ ∣ s t ) [ ∑ t = 0 T − 1 γ t π θ ( a t ∣ s t ) π θ o l d ( a t ∣ s t ) A π θ o l d ( s t , a t ) ] ≈ 1 N ∑ i = 1 N ∑ t = 0 T ( i ) − 1 γ t π θ ( a t ( i ) ∣ s t ( i ) ) π θ o l d ( a t ( i ) ∣ s t ( 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 K ( θ ; θ o l d ) = E τ ∼ ( p 0 , π θ o l d , p ) , a t ∼ π θ o l d ( ⋅ ∣ s t ) [ t = 0 ∑ T − 1 γ t π θ o l d ( a t ∣ s t ) π θ ( a t ∣ s t ) A π θ o l d ( s t , a t ) ] ≈ N 1 i = 1 ∑ N t = 0 ∑ T ( i ) − 1 γ t π θ o l d ( a t ( i ) ∣ s t ( i ) ) π θ ( a t ( i ) ∣ s t ( i ) ) A ^ t
这里 A ^ t ≈ A π θ o l d ( s t , a t ) \hat{A}_t\approx A^{\pi_{\theta_{old}}}(s_t,a_t) A ^ t ≈ A π θ o l d ( s t , a t ) 是一个 advantage estimator. 注意上面的 estimator 当 N N N 变大以及 π θ ≈ π θ o l d \pi_{\theta}\approx \pi_{\theta_{old}} π θ ≈ π θ o l d 时比较准确。
在实际实现中,由于 discount factor γ < 1 \gamma < 1 γ < 1 , 远期 advantage 对梯度的贡献会指数级衰减。为了提高效率,TRPO 中常常使用 γ = 1 \gamma = 1 γ = 1 来简化计算,但此时需要另外一种形式的约束来保证目标函数的有界性。这被称为 gamma trick :通过适当选择 γ \gamma γ 的值来平衡 bias 和 variance.
更具体地,当 γ < 1 \gamma < 1 γ < 1 时:
远期奖励对当前决策的影响被衰减,这在经济学中是合理的(现值概念).
有助于减少远期估计的方差(远期的 V π V^{\pi} V π 估计不确定性更大).
在 TRPO 的 derivation 中,γ \gamma γ 出现在目标函数的 discount 和中,使用较小的 γ \gamma γ 可以使 π θ ≈ π θ o l d \pi_{\theta} \approx \pi_{\theta_{old}} π θ ≈ π θ o l d 的近似更准确.
我们把上面的结果整理为如下算法
Algorithm: TRPO MC
Initialize policy parameters θ 0 \theta_0 θ 0 , value function parameters ϕ 0 \phi_0 ϕ 0 . Set hyperparameter δ > 0 \delta > 0 δ > 0 .
for k = 0 , 1 , … k = 0, 1, \dots k = 0 , 1 , … :
Sample N N N trajectories { τ i } \{\tau_i\} { τ i } according to ( p 0 , π θ k , p ) (p_0, \pi_{\theta_k}, p) ( p 0 , π θ k , p ) .
Solve the constrained optimization :
max θ k + 1 1 N ∑ i = 1 N ∑ t = 0 T ( i ) − 1 π θ k + 1 ( a t ( i ) ∣ s t ( i ) ) π θ k ( a t ( i ) ∣ s t ( i ) ) A ^ t s . t . max s ∈ S K L ( π θ 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} θ k + 1 max s.t. N 1 i = 1 ∑ N t = 0 ∑ T ( i ) − 1 π θ k ( a t ( i ) ∣ s t ( i ) ) π θ k + 1 ( a t ( i ) ∣ s t ( i ) ) A ^ t s ∈ S max KL ( π θ k ( ⋅ ∣ s ) ∣∣ π θ k + 1 ( ⋅ ∣ s )) ≤ δ
Compute R ^ t ( i ) = γ t ′ − t r t ′ ( i ) \hat{R}_t^{(i)} = \gamma^{t'-t} r_{t'}^{(i)} R ^ t ( i ) = γ t ′ − t r t ′ ( i ) for i = 1 , … , N i = 1, \dots, N i = 1 , … , N , t = 1 , … , T ( i ) − 1 t = 1, \dots, T^{(i)}-1 t = 1 , … , T ( i ) − 1 .
Update the value model :
min ϕ 1 N ∑ i = 1 N 1 T ( i ) ∑ t = 0 T ( i ) − 1 1 2 ( V ϕ ( s t ( 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 ϕ min N 1 i = 1 ∑ N T ( i ) 1 t = 0 ∑ T ( i ) − 1 2 1 ( V ϕ ( s t ( i ) ) − R ^ t ( i ) ) 2
关于这个优化问题的解法可以参考后续 discussion 章节
我们将目标函数和约束进行泰勒展开得到
K ( θ ; θ o l d ) ≈ g T ( θ − θ k ) K L ( θ ∣ ∣ θ k ) ≈ 1 2 ( θ − θ k ) T H ( θ − θ 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} K ( θ ; θ o l d ) KL ( θ ∣∣ θ k ) ≈ g T ( θ − θ k ) ≈ 2 1 ( θ − θ k ) T H ( θ − θ k )
这样我们的优化目标就近似为
max θ g T ( θ − θ k ) s . t . 1 2 ( θ − θ k ) T H ( θ − θ 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} θ max s.t. g T ( θ − θ k ) 2 1 ( θ − θ k ) T H ( θ − θ k ) ≤ δ
这个目标函数与 TRO (trust region methods) 一致,因此我们可以使用类似的做法来解决,我们可以得到上面问题的准确表达式
θ k + 1 = θ k + 2 δ g T H − 1 g H − 1 g \theta_{k+1} = \theta_k + \sqrt{\frac{2\delta}{g^TH^{-1}g}}H^{-1}g θ k + 1 = θ k + g T H − 1 g 2 δ H − 1 g
但是,由于我们使用了 Taylor 展开,实际上我们的更新并不一定满足 KL divergence 的约束,TRPO 对此进行了改进,即在更新时加入了 line search 来使得 θ k + 1 \theta_{k+1} θ k + 1 满足 KL divergence 约束。
θ k + 1 = θ k + α j 2 δ g T H − 1 g H − 1 g \theta_{k+1} = \theta_k + \alpha^j\sqrt{\frac{2\delta}{g^TH^{-1}g}}H^{-1}g θ k + 1 = θ k + α j g T H − 1 g 2 δ H − 1 g
其中 α ∈ ( 0 , 1 ) \alpha\in(0,1) α ∈ ( 0 , 1 ) 是 backtracking line search 的参数, j j j 是使得 θ k + 1 \theta_{k+1} θ k + 1 满足 KL divergence 约束的最小正整数 k k k .
虽然 TRPO 的理论形式非常简单,但是实际上实现起来很麻烦。其根本原因在于这个约束比较难以满足,原始论文中使用了 line search 来解决这个问题,但是当状态空间大了之后,速度会显著下降,这也是 PPO 的核心贡献之一。
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) .
在上一节中,我们介绍了 TRPO 算法,TRPO 算法是一个二阶算法,TRPO 大幅度提高了 sample efficiency, 但是其问题在于计算/优化太过复杂。因此,PPO 就通过 clip 等技巧保留了 TRPO 思想并降低了计算难度。
在 TRPO 算法中,我们的优化问题形式为:
max θ K ( θ ; θ o l d ) = E τ ∼ ( p 0 , π θ o l d , p ) , a t ′ ∼ π θ o l d ( ⋅ ∣ s t ) [ ∑ t = 0 T − 1 π θ ( a t ′ ∣ s t ) π θ o l d ( a t ′ ∣ s t ) A π θ o l d ( s t , a t ′ ) ] s . t . max s ∈ S K L ( π θ o l d ( ⋅ ∣ 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} θ max s.t. K ( θ ; θ o l d ) = E τ ∼ ( p 0 , π θ o l d , p ) , a t ′ ∼ π θ o l d ( ⋅ ∣ s t ) [ t = 0 ∑ T − 1 π θ o l d ( a t ′ ∣ s t ) π θ ( a t ′ ∣ s t ) A π θ o l d ( s t , a t ′ ) ] s ∈ S max KL ( π θ o l d ( ⋅ ∣ s ) ∣∣ π θ ( ⋅ ∣ s )) ≤ δ
TRPO 的核心思想为,当 π θ \pi_{\theta} π θ 和 π θ o l d \pi_{\theta_{old}} π θ o l d 比较接近时,我们的目标函数 K ( θ ; θ o l d ) \mathcal{K}(\theta;\theta_{old}) K ( θ ; θ o l d ) 与 policy gradient 的目标函数 J ( θ ) \mathcal{J}(\theta) J ( θ ) 梯度误差比较小。
PPO 针对这一点核心观点进行了扩展,我们可以通过另一种形式来约束 π θ \pi_{\theta} π θ 和 π θ o l d \pi_{\theta_{old}} π θ o l d 比较接近 ,即
r t ( θ ) : = π θ ( a t ′ ∣ s t ) π θ o l d ( a t ′ ∣ s t ) ≈ 1 r_t(\theta) := \frac{\pi_\theta(a_t'\mid s_t)}{\pi_{\theta_{old}}(a_t'\mid s_t)} \approx 1 r t ( θ ) := π θ o l d ( a t ′ ∣ s t ) π θ ( a t ′ ∣ s t ) ≈ 1
基于这个思想,我们可以丢弃 TRPO 的约束,直接对 r t ( θ ) r_t(\theta) r t ( θ ) 进行约束,这个约束可以使用 c l i p \mathrm{clip} clip 函数来实现,此时优化的目标函数就变成了
J ( θ ) = E ( q , a ) ∼ D , o ≤ t ∼ π θ o l d ( ⋅ ∣ q ) [ c l i p ( r t ( θ ) , 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] J ( θ ) = E ( q , a ) ∼ D , o ≤ t ∼ π θ o l d ( ⋅ ∣ q ) [ clip ( r t ( θ ) , 1 − ϵ , 1 + ϵ ) A ^ t ]
这里 ϵ > 0 \epsilon>0 ϵ > 0 是一个超参数,类似于 TRPO 中的 δ \delta δ .
c l i p ( x , ℓ , r ) = { ℓ , if x ≤ ℓ x , if ℓ < x < r r , if x ≥ r \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 ( x , ℓ , r ) = ⎩ ⎨ ⎧ ℓ , x , r , if x ≤ ℓ if ℓ < x < r if x ≥ r
现在我们已经通过 c l i p \mathrm{clip} clip 抛弃了 TRPO 复杂的约束了。我们来分析一下目标函数的性质,我们将不同的结果总结为下表
从结果中我们可以看出,只有当 r t ≈ 1 r_t\approx 1 r t ≈ 1 时,我们才会更新我们的模型,对于 A ^ t > 0 , r t < 1 − ϵ \hat{A}_t>0, r_t<1-\epsilon A ^ t > 0 , r t < 1 − ϵ 和 A ^ t < 0 , r t < 1 − ϵ \hat{A}_t<0, r_t<1-\epsilon A ^ t < 0 , r t < 1 − ϵ 这两种应该更新的情况,我们并没有更新,我们的 sample efficiency 很低。为了解决这个问题,我们在 c l i p \mathrm{clip} clip 的基础上,进一步加入一个 min \min min 函数来进行控制,这样,我们就得到了 PPO 的目标函数:
J P P O ( θ ) = E ( q , a ) ∼ D , o ≤ t ∼ π θ o l d ( ⋅ ∣ q ) [ min ( r t ( θ ) A ^ t , c l i p ( r t ( θ ) , 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] J PPO ( θ ) = E ( q , a ) ∼ D , o ≤ t ∼ π θ o l d ( ⋅ ∣ q ) [ min ( r t ( θ ) A ^ t , clip ( r t ( θ ) , 1 − ϵ , 1 + ϵ ) A ^ t ) ]
此时,我们再对目标函数进行分析,就得到
可以看到,此时 PPO 的目标函数就解决了 sample efficiency 过低的问题了。
其中
r t ( θ ) = π θ ( o t ∣ q , o < t ) π θ o l d ( o t ∣ q , o < t ) r_t(\theta) = \frac{\pi_{\theta}(o_t\mid q, o_{< t})}{\pi_{\theta_{old}}(o_t\mid q, o_{< t})} r t ( θ ) = π θ o l d ( o t ∣ q , o < t ) π θ ( o t ∣ q , o < t )
( q , a ) (q,a) ( q , a ) 是从数据集 D \mathcal{D} D 采样的 QA pair,ϵ > 0 \epsilon>0 ϵ > 0 是一个超参数,A ^ t \hat{A}_t A ^ t 是 t t t 时刻的优势估计 (advantage estimator). 给定 value function V V V 以及 reward function R R R , A ^ t \hat{A}_t A ^ t 通过计算 GAE 得到:
A ^ t G A E ( γ , λ ) = ∑ k = 0 ∞ ( γ λ ) k δ t + k \hat{A}_t^{\mathrm{GAE}(\gamma, \lambda)}=\sum_{k=0}^{\infty}(\gamma\lambda)^k\delta_{t+k} A ^ t GAE ( γ , λ ) = k = 0 ∑ ∞ ( γ λ ) k δ t + k
其中
δ k = R k + γ V ( s k + 1 ) − V ( s k ) , 0 ≤ γ , λ ≤ 1 \delta_k = R_k + \gamma V(s_{k+1})-V(s_k),\quad 0\leq \gamma,\lambda\leq 1 δ k = R k + γ V ( s k + 1 ) − V ( s k ) , 0 ≤ γ , λ ≤ 1
TODO
GRPO (Group Relative Policy Optimization) 是 DeepSeek 提出的 RL 算法,相比于 PPO,GRPO 不依赖于 value function(因此也不需要 reward model 来训练 critic),而是通过同一 prompt 下的多组输出来估计 advantage.
PPO 需要维护一个 value function V ϕ V_{\phi} V ϕ 来估计 advantage A t A_t A t :
A ^ t G A E = ∑ k = 0 ∞ ( γ λ ) k δ t + k , δ k = r k + γ V ϕ ( s k + 1 ) − V ϕ ( s k ) \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) A ^ t GAE = k = 0 ∑ ∞ ( γ λ ) k δ t + k , δ k = r k + γ V ϕ ( s k + 1 ) − V ϕ ( s k )
但是在 LLM 场景下,训练 value function 会带来额外的计算和内存开销:
需要额外训练一个与 policy model 规模相当的 critic model
需要 reward model 提供的 token-level reward(而很多 verifier 只提供 outcome-level reward)
GRPO 的核心思想:用同一 prompt 下的一组输出来估计 baseline(即群体级别的相对奖励),从而避免训练 value function.
给定 QA pair ( q , a ) (q, a) ( q , a ) , 从 π θ o l d \pi_{\theta_{old}} π θ o l d 中采样 G G G 个输出 { o i } i = 1 G \{o_i\}_{i=1}^G { o i } i = 1 G , 并获取对应的 reward { R i } i = 1 G \{R_i\}_{i=1}^G { R i } i = 1 G .
Advantage Estimation :
GRPO 使用 group-level 归一化来估计 advantage:
A ^ i , t = R i − m e a n ( { R i } i = 1 G ) s t d ( { R i } i = 1 G ) \hat{A}_{i,t} = \frac{R_i - \mathrm{mean}(\{R_i\}_{i=1}^G)}{\mathrm{std}(\{R_i\}_{i=1}^G)} A ^ i , t = std ({ R i } i = 1 G ) R i − mean ({ R i } i = 1 G )
这里 advantage A ^ i , t \hat{A}_{i,t} A ^ i , t 与时间步 t t t 无关(仅与 group 中的 reward 排名相关),这使得 GRPO 能够直接使用 outcome-level reward.
💡
直观理解:同一 prompt 下的 G G G 个 response 中,高于平均奖励的 response
获得正的 advantage, 低于平均奖励的 response 获得负的 advantage. 这与 PPO 中
A t = Q t − V t A_t = Q_t - V_t A t = Q t − V t 的思想一致:V t V_t V t 是”平均意义上的 expected return”, 起到
baseline 的作用.
GRPO Objective :
GRPO 的训练目标与 PPO 类似,但在分组上进行了归一化:
J G R P O ( θ ) = E ( q , a ) ∼ D , { o i } i = 1 G ∼ π θ o l d ( ⋅ ∣ q ) [ 1 G ∑ i = 1 G 1 ∣ o i ∣ ∑ t = 1 ∣ o i ∣ min ( r i , t ( θ ) A ^ i , t , c l i p ( r i , 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] J GRPO ( θ ) = E ( q , a ) ∼ D , { o i } i = 1 G ∼ π θ o l d ( ⋅ ∣ q ) G 1 i = 1 ∑ G ∣ o i ∣ 1 t = 1 ∑ ∣ o i ∣ min ( r i , t ( θ ) A ^ i , t , clip ( r i , t ( θ ) , 1 − ϵ , 1 + ϵ ) A ^ i , t )
其中:
r i , t ( θ ) = π θ ( o i , t ∣ q , o i , < t ) π θ o l d ( o i , t ∣ q , o i , < 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})} r i , t ( θ ) = π θ o l d ( o i , t ∣ q , o i , < t ) π θ ( o i , t ∣ q , o i , < t )
是 importance sampling ratio,衡量新旧策略在 token o i , t o_{i,t} o i , t 处的概率比.
⚠️
GRPO 中 A ^ i , t \hat{A}_{i,t} A ^ i , t 对所有 token 使用相同的 advantage 值(基于 outcome reward)。这与 PPO 使用 GAE 得到 per-token advantage 不同.
在实际的 RLHF/GRPO 训练中,通常会在 reward 中减去一个 KL 惩罚项,以防止模型偏离 reference model 太远:
R i = R t a s k ( o i ) − β K L ( π θ ∥ π r e f ) R_i = R_{\mathrm{task}}(o_i) - \beta \, \mathrm{KL}(\pi_{\theta} \| \pi_{\mathrm{ref}}) R i = R task ( o i ) − β KL ( π θ ∥ π ref )
其中 β > 0 \beta > 0 β > 0 是超参数,π r e f \pi_{\mathrm{ref}} π ref 是 reference model (通常是 SFT 后的模型).
常用的 KL 估计方式包括:
Kullback-Leibler divergence : K L ( π θ ∥ π r e f ) = E o ∼ π θ [ log π θ ( o ∣ q ) π r e f ( o ∣ q ) ] \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] KL ( π θ ∥ π ref ) = E o ∼ π θ [ log π ref ( o ∣ q ) π θ ( o ∣ q ) ]
Unbiased KL estimator : DeepSeek 使用的低方差估计,形式为:
K L ( π θ ∥ π r e f ) = π r e f ( o i , t ∣ q , o i , < t ) π θ ( o i , t ∣ q , o i , < t ) − log π r e f ( o i , t ∣ q , o i , < t ) π θ ( o i , t ∣ q , o i , < 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 KL ( π θ ∥ π ref ) = π θ ( o i , t ∣ q , o i , < t ) π ref ( o i , t ∣ q , o i , < t ) − log π θ ( o i , t ∣ q , o i , < t ) π ref ( o i , t ∣ q , o i , < t ) − 1
GRPO 的典型训练流程如下:
Algorithm: GRPO Training
Input : Dataset D \mathcal{D} D of prompts, reference model π r e f \pi_{\mathrm{ref}} π ref , group size G G G .
Initialize : policy π θ = π r e f \pi_{\theta} = \pi_{\mathrm{ref}} π θ = π ref .
for each training iteration:
Sample a batch of prompts { q } \{q\} { q } from D \mathcal{D} D .
For each prompt q q q , sample G G G responses { o i } \{o_i\} { o i } from π θ \pi_{\theta} π θ :
o i ∼ π θ ( ⋅ ∣ q ) , i = 1 , … , G o_i \sim \pi_{\theta}(\cdot \mid q), \quad i = 1, \dots, G o i ∼ π θ ( ⋅ ∣ q ) , i = 1 , … , G
Compute reward for each response:
R i = R t a s k ( o i ) − β K L ( π θ ∥ π r e f ) R_i = R_{\mathrm{task}}(o_i) - \beta \, \mathrm{KL}(\pi_{\theta} \| \pi_{\mathrm{ref}}) R i = R task ( o i ) − β KL ( π θ ∥ π ref )
Compute group-level advantage:
A ^ i = R i − m e a n ( { R } ) s t d ( { R } ) \hat{A}_i = \frac{R_i - \mathrm{mean}(\{R\})}{\mathrm{std}(\{R\})} A ^ i = std ({ R }) R i − mean ({ R })
Update π θ \pi_{\theta} π θ using GRPO objective:
L = 1 G ∑ i = 1 G 1 ∣ o i ∣ ∑ t = 1 ∣ o i ∣ min ( r i , t ( θ ) A ^ i , c l i p ( r i , 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) L = G 1 i = 1 ∑ G ∣ o i ∣ 1 t = 1 ∑ ∣ o i ∣ min ( r i , t ( θ ) A ^ i , clip ( r i , t ( θ ) , 1 − ϵ , 1 + ϵ ) A ^ i )
(Optional) Update π r e f \pi_{\mathrm{ref}} π ref periodically.
无需求解 value function : 在 LLM 场景下,action space (词表) 和 state space (上下文) 极大,训练 value function 是非常困难且昂贵的。GRPO 通过 group 比较自然避免了这一问题.
Outcome-level reward 友好 : 许多 LLM 任务(如数学推理、代码生成)只有最终的 outcome reward (答案正确与否)。GRPO 可以直接利用 outcome reward,无需 reward model 提供 token-level reward.
简单高效 : GRPO 的实现相对简单,无需维护 critic network,计算和内存开销更小.