Reinforcement Learning for Large Language models: An Overview

通常来说，state 包含了当前进行决策所需要的所有信息（环境信息，历史信息，agent 自身信息），这种情况下 MDP 的 Markov 性质成立。实际上 agent 获取的可能只是一部分信息，即 agent 获取的为 observation, 这种情况下 MDP 特化为 Partially Observable Markov Decision Process (POMDP), 此时我们的 Markov 性质就不再对 observation 成立，我们的问题也变的更加困难。

Takeaway State 表示整个环境的完整表示，而 observation 则是环境的部分表示

对于 LLM 来说，我们可以获取历史状态的所有信息，因此我们可以将 RL for LLM 视作一个 MDP 问题。

action space

action space $\mathcal{A}$ 可以是连续的，比如上下左右四个方向，也可以是连续的，比如方向盘的角度

episodic and continuing

根据任务的终止时间 $T$ ，我们可以将任务分为 episodic task 和 continuing task, 当 $T=\infty$ 时，我们的任务称为 continual task，比如游戏, 否则称之为 episodic task, 比如股票预测.

除了 POMDP 之外，MDP 也有一些其他的推广形式：

non-stationary dynamics, 即不同时间步的状态转移函数不同
pre-determined terminal time $T$ . 终止时间指定，因为 $p_{T-1}(\langle term\rangle\mid s_{T-1}, a_{T-1})=1$ 且 $p_{t+1}(\langle term\rangle\mid s_{t}, a_{t})=0$ , $\forall t\in[T-2]$ , 所以此时 MDP 为 non-stationary,
policy $\pi_t$ may be time-dependent, 一般来说仅在 non-stationary dynamics 中有必要
actions may depend on states, 即 $a_t\in\mathcal{A}(s_t)$ .

由于 $T$ 也是一个随机变量，因此 $\mathbb{E}[\sum_{t=0}^T\cdot]\neq\sum_{t=0}^T\mathbb{E}[\cdot]$ .

为了简化，我们一般会使用一个等价的，不会停止的 MDP:

the MDP never stops, 即 $T=\infty$
$s_t\in\mathcal{S}\cup\{\langle term\rangle\}$ , 这里 $\langle term\rangle$ 是一个 normal state.
absorbing state: 当 $s_t=\langle term\rangle$ 时， $r_t=0, s_{t+1}=\langle term\rangle$ , 即状态不会离开 $\langle term\rangle$ .
$\pi(\cdot\mid\langle term\rangle)$ 不产生任何影响

Model Based

RL Objective

RL 的最终目标为最大化累计奖励，称为 expected return, 这个目标基于reward hypothesis, 即所有的目标都可以被描述为最大化 expected return. 目标函数表达式如下

\max_{\pi}\quad \mathbb{E}_{s_0\sim p_0}^\pi\left[\sum_{t=0}^{T-1}r_t\right]

由于未来存在不确定性，我们会对这种不确定性施加惩罚，即时间越远，其对于当前的 reward 就越低，这其实就是经济学中的现值 (present value, PV), 因此我们实际上优化的目标函数是 expected discounted return:

\max_{\pi}\quad \mathbb{E}_{s_0\sim p_0}^\pi\left[\sum_{t=0}^{T-1}\gamma^tr_t\right]

这里 $\gamma \in (0, 1]$ 是一个超参数，我们定义

G_t = r_t +\gamma r_{t+1}+\cdots+\gamma^{T-1-t}r_{T-1}=\sum_{t'=t}^{T-1}\gamma^{t'-t}r_{t'}

$G_t$ 被称为 discounted return.

注意这里的期望 $\mathbb{E}_{s_0\sim p_0}^\pi$ 对应的随机变量为 $a_t\sim\pi(\cdot\mid s_t)$ , $(r_t,s_{t+1})\sim p(\cdot,\cdot\mid s_t,a_t)$ .

Connection with LLM

我们将 LLM 记为 $p_{\theta}$ , 这里 $\theta$ 是 LLM 的参数，将 LLM 的输入记为 $x$ , 这里 $x$ 一般就是一个 prompt, LLM 每次生成一个新的 token, 我们记为 $y_i$ , 表达式为

y_i \sim p(y\mid y_{<i}, x)

RL 和 LLM 的对应关系如下

RL	LLM
policy	LLM $p_{\theta}$
initial state	prompt $x$
action	a single output token $y_i$
state	prompt and output token $(y_{<i}, x)$
reward	verifier or reward model

注意到我们并没有对 environment 进行建模，我们可以这样进行理解：

LLM 中的 environment 仅仅是对 token 进行简单拼接
environment 承担了 reward model + 文本世界的功能，其中奖励由 reward model 给定或者人类反馈，且 environment 使用 <eos> 来标记终止条件

Conclusion

在本节中，我们介绍了 RL 的基本定义，RL 的核心在于其 Markov 性质。我们还介绍了 RL 中的核心假设，即所有的目标都可以使用 reward 量化。最后我们从 RL 的角度来对 LLM 进行建模

Bellman Equations

Value function and Q function

本节中，我们将要定义 value function, Q-function 以及这两个函数与 policy 之间的关系，这是我们介绍不同 RL 算法的基础

我们定义 Value function 如下：

V^{\pi}(s) = \mathbb{E}^{\pi}[G_0\mid s_0=s]

value function 的具体含义为：agent 从当前状态 $s$ 出发，一直遵循当前策略 $\pi$ ，最后获取到的 expected discounted return.

state-action value function, 或者 Q-value function 定义如下：

Q^{\pi}(s,a) = \mathbb{E}^{\pi}[G_0\mid s_0=s, a_0=a]

其具体含义为：agent 从当前状态 $s$ 出发，执行 action $a$ , 再遵循当前策略 $\pi$ , 最后获取到的 expected discounted return.

注意这里为了方便，我们令 $V^{\pi}(\langle term\rangle)=0$ , $Q^{\pi}(\langle term\rangle,a)=0,\forall a\in\mathcal{A}$ .

由全概率公式，我们有

\mathbb{E}[G_0\mid s_0=s] = \sum_{a\in\mathcal{A}}\mathbb{E}\left[G_0\mid s_t=s, a_0=a\right]\pi(a\mid s_0=s)

即

\boxed{V^{\pi}(s) = \mathbb{E}_{a_0\sim \pi(\cdot\mid s_0)}[Q^{\pi}(s,a)]}

一般来说，我们会使用 value function 和 Q-value function 的如下递推性质：

\boxed{ \begin{aligned}V^\pi(s) &= \mathbb{E}_{a_0\sim \pi(\cdot\mid s),\ (r_0,s_0)\sim p(\cdot,\cdot\mid s,a_0)}[r_0 + \gamma V^\pi(s_1)\mid s_0=s]\\ Q^\pi(s,a) &= \mathbb{E}_{(r,s')\sim p(\cdot,\cdot\mid s,a),\ a'\sim \pi(\cdot\mid s')}[r + \gamma Q^\pi(s',a')\mid s,a] \end{aligned}}

证明需要利用 Markov property:

\begin{aligned} V^{\pi}(s) &= \mathbb{E}^{\pi}\left[r_0 + \gamma\sum_{t=1}^{T-1}\gamma^{t-1}r_t\mid s_0=s\right]\\ &= \mathbb{E}^{\pi}[r_0+\gamma G_1\mid s_0=s]\\ &= \mathbb{E}_{a_0\sim \pi(\cdot\mid s,\ (r_0,s_0)\sim p(\cdot,\cdot\mid s,a_0))}\left[\mathbb{E}^\pi[r_0 + \gamma G_1\mid s_0,a_0,r_0,s_1]\mid s_0=s\right]\\ &= \mathbb{E}_{a_0\sim \pi(\cdot\mid s,\ (r_0,s_0)\sim p(\cdot,\cdot\mid s,a_0))}\left[r_0 + \gamma \mathbb{E}^\pi[G_1\mid s_1]\mid s_0=s\right]\\ &= \mathbb{E}_{a_0\sim \pi(\cdot\mid s,\ (r_0,s_0)\sim p(\cdot,\cdot\mid s,a_0))}\left[r_0 + \gamma V^\pi[s_1)\mid s_0=s\right] \end{aligned}

对于 $Q^\pi(s,a)$ 的推导同理。

Bellman equation Theorem

我们接下来介绍 RL 的核心：Bellman equation Theorem, 其阐述了策略 $\pi$ 与 value function 所满足的充分必要条件。

Theorem

令 $\pi$ 为一个策略，假设 $\gamma\in(0,1)$ , $|\mathcal{S}|<\infty$ 以及 $|r|\leq R<\infty, \mathrm{a.s.}$ . 那么 $\pi$ 对应的 value function $V^\pi:\mathcal{S}^+\to\mathbb{R}$ 存在，且满足 Bellman equation:

\boxed{ V^\pi(s) = \mathbb{E}_{a\sim \pi(\cdot\mid s),\ (r,s')\sim p(\cdot,\cdot\mid s,a_0))}\left[r+\gamma V(s')\mid s\right] }

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

证明

由 $V^\pi$ 定义，我们有

\left|\sum_{t=0}^{\infty}\gamma^tr_t\right|\leq \sum_{t=0}^\infty \gamma^t R=\frac{R}{1-\gamma}<\infty

因此， $\sum_{t=0}^{\infty}\gamma^tr_t$ 是一个绝对收敛的序列，从而其期望存在且有界。

对于 $V^\pi(s)=\mathbb{E}^{\pi}[G_0\mid s_0=s]$ , 我们有

\begin{aligned} V^\pi(s)&=\mathbb{E}^{\pi}\left[G_0\mid s_0=s\right]\\ &= \mathbb{E}^{\pi}\left[r_0+\gamma G_1\mid s_0=s\right]\\ &=\mathbb{E}_{a\sim\pi(\cdot\mid s), (r,s')\sim p(\cdot,\cdot\mid s,a_0)}\left[\mathbb{E}^\pi\left[r_0+\gamma G_1\mid s_0,a_0,r_0,s'\right]\right]\\ &= \mathbb{E}_{a\sim\pi(\cdot\mid s), (r,s')\sim p(\cdot,\cdot\mid s,a_0)}\left[r_0+\gamma\mathbb{E}^\pi\left[ G_1\mid s_0,a_0,r_0,s'\right]\right]\\ &= \mathbb{E}_{a\sim\pi(\cdot\mid s), (r,s')\sim p(\cdot,\cdot\mid s,a_0)}\left[r_0+\gamma\mathbb{E}^\pi\left[ G_1\mid s_1\right]\right]\\ &=\mathbb{E}_{a\sim\pi(\cdot\mid s'), (r,s')\sim p(\cdot,\cdot\mid s,a_0)}\left[r_0+V^{\pi}(s')\right] \end{aligned}

其中第 5 个等式为 Markov property.

最后，我们证明方程解的唯一性，假设存在函数 $V:\mathcal{S}^+\to\mathbb{R}$ 满足 Bellman equation. 我们定义 Bellman 算子 $\mathcal{T}^\pi$ 为

(\mathcal{T}^\pi V)(s) := \mathbb{E}_{a\sim \pi(\cdot\mid s),\ (r,s')\sim p(\cdot,\cdot\mid s,a_0))}\left[r+\gamma V(s')\mid s\right]

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

\begin{aligned} \left|(\mathcal{T}^\pi V_1)(s)-(\mathcal{T}^\pi V_2)(s)\right| &= \left|\mathbb{E}_{a\sim\pi(\cdot\mid s'), (r,s')\sim p(\cdot,\cdot\mid s,a_0)}\left[\gamma V_1^{\pi}(s')-\gamma V_2^{\pi}(s')\right]\right|\\ &\leq \gamma \mathbb{E}_{a\sim\pi(\cdot\mid s'), (r,s')\sim p(\cdot,\cdot\mid s,a_0)}\left|V_1^{\pi}(s')-V_2^{\pi}(s')\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}

上式对于任意 $s\in\mathcal{S}$ 都成立，因此

\left\|\mathcal{T}^\pi V_1-\mathcal{T}^\pi V_2\right\|_\infty \leq \gamma \left\|V_1-V_2\right\|_\infty

由于 $\gamma < 1$ , 从而 $\mathcal{T}^\pi$ 是一个 contraction mapping.

根据不动点定理，已知 $V^\pi$ 是一个不动点，而不动点唯一，因此我们有 $V=V^\pi$ . $\blacksquare$

同理，我们可以推导出关于 Q-function 的 Bellman equation:

Theorem

令 $\pi$ 为一个策略，假设 $\gamma\in(0,1)$ , $|\mathcal{S}|<\infty$ $|\mathcal{A}|<\infty$ 以及 $|r|\leq R<\infty, \mathrm{a.s.}$ . 那么 $\pi$ 对应的 Q-function $Q^\pi:\mathcal{S}^+\times\mathcal{A}\to\mathbb{R}$ 存在，且满足 Bellman equation:

\boxed{ Q^\pi(s, a) = \mathbb{E}_{(r,s')\sim p(\cdot,\cdot\mid s,a),\ a'\sim \pi(\cdot\mid s')}\left[r+\gamma Q(s', a')\mid s, a\right] }

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

证明

证明与 value function 的证明基本类似，我们这里略过。

Bellman Optimal Equation

前面介绍了针对一般策略的 Bellman equation, 特别地，对于我们的目标最优策略，我们也可以而推导出对应的 Bellman equation.

首先我们定义最优策略如下

Definition

如果策略 $\pi^*$ 满足

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

则我们称策略 $\pi^*$ 是 optimal policy. 对应的 $V^{\pi^*}$ 和 $Q^{\pi^*}$ 分别称之为 optimal value function 以及 optimal Q-function, 我们简记为 $V^*=V^{\pi^*}$ , $Q^*=Q^{\pi^*}$ . 注意 optimal policy 与状态 $s$ 无关，并且 optimal policy 是不唯一的，但是所有的 optimal policy 对应的 optimal value function 是同一个。 $V^*$ 与 $Q^*$ 之间存在如下关系

V^*(s) = \max_{a\sim\pi^*} Q^*(s, a)

证明

我们先证明左边小于右边，再证明右边小于左边。

V^*(s) =\mathbb{E}_{a_0\sim \pi^*(\cdot\mid s_0)}[Q^*(s,a)]\leq \mathbb{E}_{a_0\sim \pi^*(\cdot\mid s_0)}[\max_{a\sim\pi^*}Q^*(s,a)] = \max_{a\sim\pi^*} Q^*(s, a)

其次，令 $a^*\in\arg\max_{a}Q^*(s,a)$ , 令策略 $\pi'$ 为确定性策略 $\pi'(a^*\mid s)=1$ , 则

\max_aQ^*(s,a^*) = Q^{\pi'}(s,a^*)=V^{\pi'}(s)\leq V^*(s)

这样，我们就证明了上面的等式。 $\blacksquare$

Theorem

假设 $\gamma\in(0,1)$ , $|\mathcal{S}|<\infty$ $|\mathcal{A}|<\infty$ 以及 $|r|\leq R<\infty, \mathrm{a.s.}$ . 那么 optimal value function $V^\pi:\mathcal{S}^+\to\mathbb{R}$ 存在，且满足 Bellman optimality equation:

\boxed{ V^*(s) = \max_{a\in\mathcal{A}}\mathbb{E}_{\ (r,s')\sim p(\cdot,\cdot\mid s,a_0)}\left[r+\gamma V^*(s')\mid s, a\right] }

反之，如果存在函数 $V:\mathcal{S}^+\to\mathbb{R}$ 满足 Bellman optimality equation, 则 $V=V^*$ , 最终

\pi^*(s) = \arg\max_{a\in\mathcal{A}}\mathbb{E}_{\ (r,s')\sim p(\cdot,\cdot\mid s,a_0)}\left[r+\gamma V^*(s')\mid s, a\right]=\arg\max_{a\in\mathcal{A}} Q^*(s,a)

是一个 optimal deterministic policy.

证明

我们首先证明存在性，我们定义

(\mathcal{T}^* V)(s) := \max_{a\in\mathcal{A}}\mathbb{E}_{(r,s')\sim p(\cdot,\cdot\mid s,a_0))}\left[r+\gamma V(s')\mid s, a\right]

我们证明 $\mathcal{T}^*V$ 是一个 contraction mapping.

\begin{aligned} \left|(\mathcal{T}^* V_1)(s)-(\mathcal{T}^* V_2)(s)\right| &= \left|\max_{a\in\mathcal{A}}\mathbb{E}_{(r,s')\sim p(\cdot,\cdot\mid s,a_0))}\left[r+\gamma V_1(s')\mid s, a\right]- \max_{a\in\mathcal{A}}\mathbb{E}_{(r,s')\sim p(\cdot,\cdot\mid s,a_0))}\left[r+\gamma V_2(s')\mid s, a\right]\right|\\ &\leq \max_{a\in\mathcal{A}}\left|\mathbb{E}_{(r,s')\sim p(\cdot,\cdot\mid s,a_0))}\left[r+\gamma V_1(s')\mid s, a\right]- \mathbb{E}_{(r,s')\sim p(\cdot,\cdot\mid s,a_0))}\left[r+\gamma V_2(s')\mid s, a\right]\right|\\ &= \max_{a\in\mathcal{A}}\left|\mathbb{E}_{(r,s')\sim p(\cdot,\cdot\mid s,a_0))}\left[\gamma V_1(s')-V_2(s')\mid s, a\right]- \mathbb{E}_{(r,s')\sim p(\cdot,\cdot\mid s,a_0))}\left[r+\gamma V_2(s')\mid s, a\right]\right|\\ &\leq \gamma \max_{a\in\mathcal{A}}\left|\mathbb{E}_{a\sim\pi(\cdot\mid s'), (r,s')\sim p(\cdot,\cdot\mid s,a_0)}\left[V_1^{\pi}(s')-V_2^{\pi}(s')\mid s,a\right]\right|\\ &\leq \gamma \max_{a\in\mathcal{A}}\mathbb{E}_{a\sim\pi(\cdot\mid s'), (r,s')\sim p(\cdot,\cdot\mid s,a_0)}\left[\left|V_1^{\pi}(s')-V_2^{\pi}(s')\right|\mid s,a\right]\\ &\leq \gamma \max_{a\in\mathcal{A}}\max_{s''\in\mathcal{S}}\left|V_1(s'')-V_2(s'')\right|\\ &= \gamma \left\|V_1-V_2\right\|_\infty \end{aligned}

这里第一个不等式我们使用了结论

\left|\max_s v(s)-\max_s u(s)\right|\leq \max_s|u(s)-v(s)|

上式对于任意 $s\in\mathcal{S}$ 都成立，因此

\left\|\mathcal{T}^* V_1-\mathcal{T}^* V_2\right\|_\infty \leq \gamma \left\|V_1-V_2\right\|_\infty

由于 $\gamma < 1$ , 从而 $\mathcal{T}^*$ 是一个 contraction mapping.

根据不动点定理， $\mathcal{T}^*$ 存在一个不动点，我们将其记为 $V^*$ .

接下来，我们定义 $\pi^*$ 为

\pi^*(s) = \arg\max_{a\in\mathcal{A}}\mathbb{E}_{\ (r,s')\sim p(\cdot,\cdot\mid s,a_0)}\left[r+\gamma V^*(s')\mid s, a\right]

此时，我们有

\begin{aligned} V^*(s) &= \max_{a\in\mathcal{A}}\mathbb{E}_{\ (r,s')\sim p(\cdot,\cdot\mid s,a_0)}\left[r+\gamma V^*(s')\mid s, a\right]\\ &= \mathbb{E}_{a\sim\pi^*(s),\ (r,s')\sim p(\cdot,\cdot\mid s,a_0)}\left[r+\gamma V^*(s')\mid s\right]\\ &= \mathbb{E}^{\pi^*}\left[r_0+\gamma V^*(s_1)\mid s_0=s\right]\\ &= \mathbb{E}^{\pi^*}\left[r_0+\gamma \mathbb{E}^{\pi^*}\left[r_1+\gamma V^*(s_2)\mid s_0=s_1\right]\\\mid s_0=s\right]\\ &= \mathbb{E}^{\pi^*}\left[r_0+\gamma \mathbb{E}^{\pi^*}\left[r_1+\gamma V^*(s_2)\mid s_0=s_1, r_0\right]\\\mid s_0=s\right]\\ &= \mathbb{E}^{\pi^*}\left[\mathbb{E}^{\pi^*}\left[r_0+\gamma (r_1+\gamma V^*(s_2))\mid s_0=s_1, r_0\right]\\\mid s_0=s\right]\\ &= \mathbb{E}^{\pi^*}\left[r_0+\gamma r_1+\gamma^2 V^*(s_2)\mid s_0=s\right]\\ &= \mathbb{E}^{\pi^*}\left[r_0+\gamma r_1+\gamma^2r_2+\cdots\mid s_0=s\right]\\ &= V^{\pi^*}(s) \end{aligned}

即 $V^*=V^{\pi^*}$ .

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

V^\pi \leq V^*= \mathcal{T}^*(V^\pi)\leq ( \mathcal{T}^*)^2(V^\pi) \leq\cdots\leq (\mathcal{T}^*)^k(V^\pi) \to V^*=V^{\pi^*}

这里第一个等式是 Bellman equation, 第一个不等式是由于下面的 Lemma,最后的极限使用了不动点定理。因此我们有 $V^\pi\leq V^*$ , 从而 $\pi^*$ 是最优策略，且 $V^*$ 是对应的 optimal value function. $\blacksquare$

Lemma

令 $\pi$ 为一个策略， $\mathcal{T}^\pi$ 和 $\mathcal{T}^*$ 分别是 Bellman operator 和 Bellman optimality operator. 对任意 $V:\mathcal{S}^+\to\mathbb{R}$ , 我们有

\mathcal{T}^\pi(V) \leq \mathcal{T}^*(V)

进一步，对任意 $U:\mathcal{S}^+\to\mathbb{R}$ , $V:\mathcal{S}^+\to\mathbb{R}$ , 如果 $U\leq V$ , 则

\mathcal{T}^*(U) \leq \mathcal{T}^*(V)

证明

\begin{aligned} \mathcal{T}^\pi(V)(s) &= \mathbb{E}_{a\sim \pi(\cdot\mid s),\ (r,s')\sim p(\cdot,\cdot\mid s,a_0))}\left[r+\gamma V(s')\mid s\right]\\ &\leq \mathbb{E}_{a\sim \pi(\cdot\mid s)}\left[\max_{a'} \mathbb{E}_{(r,s')\sim p(\cdot,\cdot\mid s,a')}\left[r+\gamma V(s')\mid s\right]\right]\\ &= \max_{a'} \mathbb{E}_{(r,s')\sim p(\cdot,\cdot\mid s,a')}\left[r+\gamma V(s')\mid s\right]\\ &= \mathcal{T}^*(V) \end{aligned}

其次，令 $a^*\in\arg\max_{a}\mathbb{E}[r+\gamma U(s')]$ , 那么

\begin{aligned} \mathcal{T}^*(U)(s) &= \mathbb{E}_{(r,s')\sim p(\cdot,\cdot\mid s,a'))}\left[r+\gamma U(s')\mid s\right]\\ &\leq \mathbb{E}_{(r,s')\sim p(\cdot,\cdot\mid s,a'))}\left[r+\gamma V(s')\mid s\right]\\ &\leq \max_a\mathbb{E}_{(r,s')\sim p(\cdot,\cdot\mid s,a'))}\left[r+\gamma V(s')\mid s\right]\\ &= \mathcal{T}^*(V) \end{aligned}

证毕。 $\blacksquare$

对于 Q-function, 我们也有类似结论

Theorem

假设 $\gamma\in(0,1)$ , $|\mathcal{S}|<\infty$ $|\mathcal{A}|<\infty$ 以及 $|r|\leq R<\infty, \mathrm{a.s.}$ . 那么 optimal Q-function $Q^*:\mathcal{S}^+\times\mathcal{A}\to\mathbb{R}$ 存在，且满足 Bellman optimality equation:

\boxed{ Q^*(s, a) = \mathbb{E}_{(r,s')\sim p(\cdot,\cdot\mid s,a)}\left[r+\gamma \max_{a'\in\mathcal{A}}Q^*(s', a')\mid s, a\right] }

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

证明

证明与 value function 类似，我们这里略过。

Value Iteration

首先第一类算法就是根据 Bellman optimality equation 构造得到的 iteration methods. 由前面的分析，我们知道，函数

(\mathcal{T}^* V)(s) := \max_{a\in\mathcal{A}}\mathbb{E}_{(r,s')\sim p(\cdot,\cdot\mid s,a_0))}\left[r+\gamma V(s')\mid s, a\right]

是一个 contraction mapping, 因此我们可以构造函数序列

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

由不动点定理，我们知道这个序列收敛到 $V^*$ . 因此，value iteration 算法定义如下

\begin{aligned} V^{k+1}&=\mathcal{T}^* V^k\\ V^{k+1}(s) &= \max_{a\in\mathcal{A}}\mathbb{E}_{(r,s')\sim p(\cdot,\cdot\mid s,a_0))}\left[r+\gamma V^k(s')\mid s, a\right]\\ \pi^{k+1}(s) &= \arg\max_{a\in\mathcal{A}}\mathbb{E}_{\ (r,s')\sim p(\cdot,\cdot\mid s,a_0)}\left[r+\gamma V^k(s')\mid s, a\right] \end{aligned}

类似地，我们还可以得到 Q-value iteration algorithm

\begin{aligned} Q^{k+1} &=\mathcal{T}^*Q^k\\ Q^{k+1}(s) &=\mathbb{E}_{(r,s')\sim p(\cdot,\cdot\mid s,a)}\left[r+\gamma \max_{a'\in\mathcal{A}}Q^k(s', a')\mid s, a\right]\\ \pi^{k+1}(s) &= \arg\max_{a\in\mathcal{A}}Q^k(s,a) \end{aligned}

Policy Evaluation

这一节我们讨论如何评估 (evaluate) 一个 policy. 要评估一个 policy $\pi$ , 我们需要求出对应的 value function $V^\pi$ .

Monte Carlo

注意到

V^\pi(s) = \mathbb{E}^\pi\left[\sum_{t=0}^{T-1}\gamma^tr_t\mid s_0=s\right]

我们可以利用 Monte Carlo (MC) 方法来进行估计，即从 $s_0=s$ 出发，独立随机采样 $N$ 条轨迹

\{s, a_0^{(i)},r_0^{(i)},\dots,a_{T^{(i)}-1}^{(i)},r_{T^{(i)}-1}^{(i)}, s^{(i)}_{T^{(i)}}\}, i=1,\dots,N

然后我们使用样本平均来近似期望

V^\pi(s) = \mathbb{E}^\pi\left[\sum_{t=0}^{T-1}\gamma^tr_t\mid s_0=s\right]\approx \frac{1}{N}\sum_{i=1}^N\sum_{t=0}^{T^{(i)}-1}\gamma^tr_t^{(i)}

如果 $\mathcal{S}$ 比较小，我们可以遍历 $s\in\mathcal{S}$ 来估计 $V^\pi$ .

当 $\mathcal{S}$ 非常大或者连续时，我们需要使用神经网络 $V_\phi$ 来近似 $V^\pi$ ,

\min_{\phi} \mathcal{L}(\phi) =\mathbb{E}_{s\sim p_0}\left[\frac12\left(V_\phi(s)-V^\pi(s)\right)^2\right]

对上面的目标函数求导得到

\begin{align} \nabla_\phi \mathcal{L}(\phi) &= \mathbb{E}_{s\sim p_0}\left[\left(V_\phi(s)-V^\pi(s)\right)\nabla_\phi V_\phi(s)\right]\\ &= \mathbb{E}_{s\sim p_0}\left[\left(V_\phi(s)- \mathbb{E}^\pi\left[\sum_{t=0}^{T-1}\gamma^tr_t\mid s_0=s\right]\right)\nabla_\phi V_\phi(s)\right]\\ &=\mathbb{E}_{s\sim p_0}\left[\mathbb{E}^\pi\left[\left(V_\phi(s)- \sum_{t=0}^{T-1}\gamma^tr_t\right)\nabla_\phi V_\phi(s)\mid s_0=s\right]\right]\\ &= \mathbb{E}_{s\sim p_0}^\pi\left[\left(V_\phi(s)- \sum_{t=0}^{T-1}\gamma^tr_t\right)\nabla_\phi V_\phi(s)\mid s_0=s\right]\\ \end{align}

这样，我们可以结合 SGD 与 MC 来近似 $V^\pi$

Algorithm: Value function approximation with MC

while not converged:

$s_0\sim p_0$ , set $t=0$ .
while $s_t\neq \langle term\rangle$ $s_{t} \neq = ⟨ t er m ⟩$ :
- $a_t\sim \pi(\cdot\mid s_t)$
- $(r_t,s_{t+1})\sim p(\cdot,\cdot\mid s_t,a_t)$
- $t\gets t+1$
Set $T=t$ .
$g=\left(V_\phi(s_0)- \sum_{t=0}^{T-1}\gamma^tr_t\right)\nabla_\phi V_\phi(s_0)$ .
Update $\phi$ using $g$ with an optimizer.

Temporal Difference

Temporal difference (TD) learning 通过相邻两步的转换关系来近似函数。注意到，

V^\pi(s) =\mathbb{E}^\pi[r_0 + \gamma V^\pi(s_1)\mid s_0=s]

与前面的方法类似，我们先采样然后使用 MC 来进行估计

V^\pi(s) =\mathbb{E}^\pi[r_0 + \gamma V^\pi(s_1)\mid s_0=s]\approx \frac{1}N\sum_{i=1}^N\left(r_0^{(i)} + \gamma V^\pi(s_1^{(i)})\right)

同样的，当 $\mathcal{S}$ 比较大或者连续时，我们构造损失函数

\min_{\phi} \mathcal{L}(\phi) =\mathbb{E}_{s\sim p_0}\left[\frac12\left(V_\phi(s)-V^\pi(s)\right)^2\right]

对应的梯度为

\nabla_\phi \mathcal{L}(\phi) =\mathbb{E}_{s\sim p_0}^\pi\left[\left(V_\phi(s)- r_0-\gamma V^\pi(s_1)\right)\nabla_\phi V_\phi(s)\mid s_0=s\right]

但是，这里存在的问题在于，我们使用了 $V^\pi(s_1)$ , 而这个值是未知的，因此，一个做法是使用当前的 value function $V_\phi$ 来进行代替，即

\nabla_\phi \mathcal{L}(\phi) \approx \mathbb{E}_{s\sim p_0}^\pi\left[\left(V_\phi(s)- r_0-\gamma V_\phi(s_1)\right)\nabla_\phi V_\phi(s)\mid s_0=s\right]

这样，我们结合 TD 和 GSD 的算法就变成了

Algorithm: Value function approximation with TD (not correct)

while not converged:

$s_0\sim p_0$ , $a_0\sim\pi(\cdot\mid s_0)$ , $(r_0,s_1)\sim p(\cdot, \cdot\mid s_0,a_0)$ .
$y=\begin{cases} r_0+\gamma V_\phi(s_1),& \text{if } s_1\neq \langle term\rangle\\ r_0 & \text{otherwise} \end{cases}$
$g=(V_{\phi}(s_0)- y)\nabla_\phi V_{\phi}(s_0)$ .
Update $\phi$ using $g$ with an optimizer.

接下来，我们需要关注如何计算 $g$ , 直接通过 $g=\left(V_\phi(s)- r_0-\gamma V_\phi(s_1)\right)\nabla_\phi V_\phi(s)$ 来进行计算，这在数学上是没有问题的，但是从自动微分的角度不对，比如 Pytorch 对应的实现应该是

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

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

\nabla_\phi\left[\frac12\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)-\nabla_\phi V_\phi(s'))

这显然与 $g$ 不相等。

为了解决这个问题，我们需要使用 stop-gradient 的技巧来避免 $V_\phi(s')$ 参与反向传播，这个时候，我们的目标函数就变成了

\frac12\left(V_\phi(s)-(r+\gamma\ \mathrm{sg}[V_\phi(s')])\right)^2

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

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

这样其梯度就是

\nabla_\phi\left[\frac12\left(V_\phi(s)-(r+\gamma\ \mathrm{sg}[V_\phi(s')])\right)^2\right]=(V_\phi(s)-(r+\gamma V_\phi(s'))\nabla_\phi V_\phi(s)=g

对应的 python 代码为

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

最终，我们的算法实现为

Algorithm: Value function approximation with TD

while not converged:

$s_0\sim p_0$ , $a_0\sim\pi(\cdot\mid s_0)$ , $(r_0,s_1)\sim p(\cdot, \cdot\mid s_0,a_0)$ .
$y=\begin{cases} r_0+\gamma\, \mathrm{sg}[V_\phi(s_1)],& \text{if } s_1\neq \langle term\rangle\\ r_0 & \text{otherwise} \end{cases}$
$g=(V_{\phi}(s_0)- y)\nabla_\phi V_{\phi}(s_0)$ .
Update $\phi$ using $g$ with an optimizer.

上述这种做法其实是 semi-gradient methods, 这类方法在参数更新时，只计算部分梯度的方法。虽然这种方法牺牲了严格的梯度下降性质，但是其能够保证更好的表现。

K-step TD

我们可以进一步推广 TD 到多步的场景，注意到

\begin{align} V^\pi(s) &=\mathbb{E}^\pi[r_0 + \gamma V^\pi(s_1)\mid s_0=s]\\ &=\mathbb{E}^\pi[r_0 + \gamma \mathbb{E}^\pi[r_1 + \gamma V^\pi(s_1)\mid s_1]\mid s_0=s]\\ &= \mathbb{E}^\pi[\mathbb{E}^\pi[r_0 + \gamma r_1 + \gamma^2 V^\pi(s_2) \mid s_1] \mid s_0=s] \end{align}

由重期望定理（Law of Total Expectation）：

\mathbb{E}^\pi[\mathbb{E}^\pi[X \mid s_1] \mid s_0=s] = \mathbb{E}^\pi[X \mid s_0=s]

因此：

V^\pi(s) = \mathbb{E}^\pi[r_0 + \gamma r_1 + \gamma^2 V^\pi(s_2) \mid s_0=s]

重复这个过程 k 次,得到 k 步展开：

\begin{align} V^\pi(s) &= \mathbb{E}^\pi\left[\sum_{i=0}^{k-1} \gamma^i r_i + \gamma^k V^\pi(s_k) \mid s_0=s\right] \end{align}

现在，我们可以基于 k-step transition 构建目标函数

\min_{\phi} \mathcal{L}(\phi) =\mathbb{E}_{s\sim p_0}\left[\frac12\left(V_\phi(s)-\left(\sum_{i=0}^{k-1} \gamma^i r_i + \gamma^k V^\pi(s_k)\right)\right)^2\right]

使用前面分析的方法，我们就可以写出类似的算法

Algorithm: Value function approximation with k-step TD

while not converged:

$s_0\sim p_0$ , set $t=0$ , $T=\infty$ .
for $t=0,\ldots,k-1$ $t = 0, \dots, k - 1$ :
- $a_t\sim\pi(\cdot\mid s_t)$
- $(r_t,s_{t+1})\sim p(\cdot, \cdot\mid s_t,a_t)$
$y=\sum_{t=0}^{k-1}\gamma^tr_t+\gamma^k \mathrm{sg}[V_\phi(s_k)]$ .
$g=(V_{\phi}(s_0)- y)\nabla_\phi V_{\phi}(s_0)$ .
Update $\phi$ using $g$ with an optimizer.

一般来说，我们会令 $k=5$ .

MC v.s. TD

TD 本质上一种 bootstrapping, 当需要使用 $V^\pi$ 时，我们使用 $V_\phi$ 来代替。两者对比如下

MC 需要完整的一次采样，但是 TD 可以通过 bootstrap 提高采样效率。
当 $V_\phi$ 与 $V^\pi$ 差距较大时，boostraping 会导致训练不稳定，而 MC 则相对稳定

Q-function

对于 Q-function, 我们也可以设计类似的算法。

对于 MC, 我们有

Q^{\pi}(s,a) = \mathbb{E}^{\pi}\left[\sum_{t=0}^{T-1}\gamma^tr_t\mid s_0=s, a_0=a\right]\approx \frac1N\sum_{i=1}^N\sum_{t=0}^{T^{(i)}-1}\gamma^tr_t^{(i)}

当使用模型来近似时，我们的目标函数为

\min_{\phi} \mathcal{L}(\phi) =\mathbb{E}_{s\sim p_0, a\sim\pi(\cdot\mid s)}\left[\frac12\left(Q_\phi(s,a)-Q^\pi(s,a)\right)^2\right]

对应的梯度

\nabla_\phi \mathcal{L}(\phi)= \mathbb{E}_{s\sim p_0}^\pi\left[\left(Q_\phi(s_0,a_0)-\sum_{t=0}^{T-1}\gamma^tr_t\right)\nabla_\phi Q_\phi(s_0,a_0)\right]

下面是对应的算法

Algorithm: Q function approximation with MC

while not converged:

$s_0\sim p_0$ , $a_0\sim\pi(\cdot\mid s_0)$ , set $t=0$ .
while $s_t\neq \langle term\rangle$ $s_{t} \neq = ⟨ t er m ⟩$ :
- $a_t\sim \pi(\cdot\mid s_t)$
- $(r_t,s_{t+1})\sim p(\cdot,\cdot\mid s_t,a_t)$
- $t\gets t+1$
Set $T=t$ .
$g=\left(Q_\phi(s_0,a_0)-\sum_{t=0}^{T-1}\gamma^tr_t\right)\nabla_\phi Q_\phi(s_0,a_0)$ .
Update $\phi$ using $g$ with an optimizer.

1-step Q-function TD learning

Algorithm: Q function approximation with TD

while not converged:

$s_0\sim p_0$ , $a_0\sim\pi(\cdot\mid s_0)$ , $(r_0,s_1)\sim p(\cdot, \cdot\mid s_0,a_0)$ , $a_1\sim\pi(\cdot\mid s_1)$ .
$y=\begin{cases} r_0+\gamma\, \mathrm{sg}[Q_\phi(s_1,a_1)],& \text{if } s_1\neq \langle term\rangle\\ r_0 & \text{otherwise} \end{cases}$
$g=\left(Q_\phi(s_0,a_0)-y\right)\nabla_\phi Q_\phi(s_0,a_0)$ .
Update $\phi$ using $g$ with an optimizer.

k-step Q-function TD learning

Algorithm: Q function approximation with k-step TD

while not converged:

$s_0\sim p_0$ , set $t=0$ , $T=\infty$ .
for $t=0,\ldots,k-1$ $t = 0, \dots, k - 1$ :
- $a_t\sim\pi(\cdot\mid s_t)$
- $(r_t,s_{t+1})\sim p(\cdot, \cdot\mid s_t,a_t)$
$y=\sum_{t=0}^{k-1}\gamma^tr_t+\gamma^k \mathrm{sg}[Q_\phi(s_k,a_k)]$
$g=(Q_{\phi}(s_0,a_0)- y)\nabla_\phi Q_{\phi}(s_0,a_0)$
Update $\phi$ using $g$ with an optimizer.

Policy Iteration Methods

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

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

policy iteration 算法的正确性如下

Theorem

考虑上面的 policy iteration 算法，假设 $\gamma\in(0,1)$ , $|\mathcal{S}|<\infty$ $|\mathcal{A}|<\infty$ 以及 $|r|\leq R<\infty, \mathrm{a.s.}$ . 那么

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

并且，存在 $K\in\mathbb{N}$ 满足 $V^{\pi^k}=V^*$ . 也就是说，有限步之后，我们就可以得到最优策略。

证明

由于 $V^{\pi^{k+1}}, V^{\pi^k}$ 是 value function, 因此它们满足 Bellman equation

\begin{align} V^{\pi^{k+1}} &= \mathbb{E}_{a\sim \pi(\cdot\mid s),\ (r,s')\sim p(\cdot,\cdot\mid s,a_0))}\left[r^{k+1}+\gamma V^{\pi^{k+1}}(s')\mid s\right]\\ V^{\pi^{k}} &= \mathbb{E}_{a\sim \pi(\cdot\mid s),\ (r,s')\sim p(\cdot,\cdot\mid s,a_0))}\left[r^k+\gamma V^{\pi^{k}}(s')\mid s\right] \end{align}

因为

\pi^{k+1}(s) = \arg\max_{a\in\mathcal{A}}\mathbb{E}_{a\sim \pi(\cdot\mid s),\ (r,s')\sim p(\cdot,\cdot\mid s,a_0))}\left[r+\gamma V^{\pi^k}(s')\mid s,a\right]

我们有

r+\gamma V^{\pi^{k+1}}(s') \geq r+\gamma V^{\pi^k}(s')

从而 $V^{\pi^{k+1}}\geq V^{\pi^k}$ . $\blacksquare$

Experiments

Value Based Methods

value-based methods 的 motivation 就是给当前策略 $\pi$ 一个引导，让 agent 学会执行达到更有价值的 states 的策略。

TD Learning

Function Form

当 state space 太大或者连续的 state space 时，我们前面的方法就不适用了，我们需要使用模型 $V_{\phi}$ 来近似 $V^{\pi}$ , 我们将目标函数设置为

J(\phi) = \mathbb{E}_{s}\left[\frac12\left(V^{\pi}(s)-V_{\phi}(s)\right)^2\right]

Remark 这里我们假设 $s\in\mathcal{S}$ 通过均匀采样得到。

我们可以使用随机梯度下降法来进行求解，目标函数的梯度为

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

从而随机梯度下降法的形式为

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

但是现在我们需要计算 $V^{\pi}(s)$ , 我们有两种方法来近似：

第一种方法是使用 MC 来进行近似，即

V^{\pi}(s)= \mathbb{E}^{\pi}[G_0\mid s_0=s]\approx G_0

第二种方法是使用 TD learning 来近似

V^{\pi}(s) \approx r_t + \gamma V_{\phi}(s_{t+1})

DQN

与上一节一样，DQA 实际上就是针对 optimal policy 应用 TD learning，我们目标函数为

Policy Based Methods

policy-based methods 的 motivation 很好理解，我们想要找到一个 optimal policy, 那我们就采集一些数据，通过一些方法来不停改进当前策略 $\pi$ , 直到当前策略达到最优。

根据采样数据所使用 policy 的不同，policy-based methods 一般会被分为两类：

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

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

on-policy 和 off-policy 对比

	on-policy	off-policy
examples	PPO, TRPO	DQN, DPO, DDPG
sampled from	current policy	arbitrary policy
exploration	yes	no
infra efficiency	low (needs to generate new sample)	high
training stability	unstable	more stable
compute cost	high	low
best for	exploration is needed	reused past experiences

Policy Gradient Theorem

我们的目标仍然是最大化累计奖励，目标函数为

\max_{\pi}\mathcal{J}(\theta)=\mathbb{E}_{s_0\sim p_0}\left[V^{\pi_\theta}(s_0)\right]=\mathbb{E}_{s_0\sim p_0}^{\pi_\theta}\left[\sum_{t=0}^{T-1}\gamma^tr_t\right]

我们首先将采集到的轨迹数据 $\tau$ 进行展开

\tau = (s_0,a_0,r_0,\dots,s_{T-1}, a_{T-1}, r_{T-1}, s_T)

$\tau$ 的概率分布可以写为

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)

我们使用 cumulative return 的定义，就得到

\mathcal{J}(\theta)=\mathbb{E}_{s_0\sim p_0}^\pi\left[\sum_{t=0}^{T-1}\gamma^tr_t\right] = \mathbb{E}_{s_0\sim p_0,a_t\sim\pi_\theta(\cdot\mid s_t), (r_t,s_{t+1})\sim p(\cdot,\cdot\mid s_t,a_t)}^\pi\left[\sum_{t=0}^{T-1}\gamma^tr_t\right] = \mathbb{E}_{\tau\sim(p_0,\pi_\theta, p)}\left[G_0(\tau)\right]

此时，目标函数的梯度为

\begin{align} \nabla_\theta \mathcal{J}(\theta)&=\nabla_\theta\mathbb{E}_{\tau\sim(p_0,\pi_\theta, p)}\left[G_0(\tau)\right]\\ &= \sum_\tau G_0(\tau)\nabla_\theta p_\theta(\tau)\\ &= \sum_\tau p_\theta(\tau)G_0(\tau) \nabla_\theta \log p_\theta(\tau)\\ &= \mathbb{E}_{\tau\sim(p_0,\pi_\theta, p)}\left[G_0(\tau)\nabla_\theta \log p_\theta(\tau)\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)G_0(\tau)\right] \end{align}

这里我们使用了如下不等式

\nabla_\theta p_\theta(\tau)=p_\theta(\tau)\nabla_\theta \log p_\theta(\tau)

以及

\nabla_\theta \log 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] = \sum_{t=0}^{T-1}\log\pi_\theta(a_t\mid s_t)

一些教材上还有另一种表达形式，使用了 Q-function, 即

\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(s_t,a_t)\right]

我们可以对其进行证明，证明需要用到 Expectation 中的全概率公式。注意到

\begin{align} \nabla_\theta \mathcal{J}(\theta)&=\mathbb{E}_{\tau\sim(p_0,\pi_\theta, p)}\left[\sum_{t=0}^{T-1}G_0(\tau)\nabla_\theta \log \pi_\theta(a_t\mid s_t)\right]\\ &= \sum_{t=0}^{T-1}\mathbb{E}_{\tau\sim(p_0,\pi_\theta, p)}\left[G_0(\tau)\nabla_\theta \log\pi_\theta(a_t\mid s_t)\right]\\ &=\sum_{t=0}^{T-1}\mathbb{E}_{\tau_{:t}\sim(p_0,\pi_\theta, p)}\left[\mathbb{E}_{\tau_{:t}\sim(p_0,\pi_\theta, p)}\left[G_0(\tau)\nabla_\theta \log\pi_\theta(a_t\mid s_t)\mid \tau_{:t}\right]\right]\\ &=\sum_{t=0}^{T-1}\mathbb{E}_{\tau_{:t}\sim(p_0,\pi_\theta, p)}\left[\nabla_\theta \log\pi_\theta(a_t\mid s_t)\mathbb{E}_{\tau_{:t}\sim(p_0,\pi_\theta, p)}\left[G_0(\tau)\mid \tau_{:t}\right]\right]\\ &=\sum_{t=0}^{T-1}\mathbb{E}_{\tau_{:t}\sim(p_0,\pi_\theta, p)}\left[\nabla_\theta \log\pi_\theta(a_t\mid s_t)\mathbb{E}_{\tau_{:t}\sim(p_0,\pi_\theta, p)}\left[G_0(\tau)\mid s_t,a_t\right]\right]\\ &=\sum_{t=0}^{T-1}\mathbb{E}_{\tau_{:t}\sim(p_0,\pi_\theta, p)}\left[\nabla_\theta \log\pi_\theta(a_t\mid s_t)Q^\pi(s_t,a_t)\right] \end{align}

这里第三个等式使用了全概率公式，第四个等式是因为 $\nabla_\theta \log\pi_\theta(a_t\mid s_t)$ 相对于 $\tau_{:t}$ 来说是一个常数，第五个等式是因为 Markov property, 最后一个等式使用了 Q-function 的定义。

Vanilla Policy Gradient Methods

令

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

则 $\hat{g}$ 是 $\nabla_\theta \mathcal{J}(\theta)$ 的一个无偏估计。

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

	\begin{algorithm}
	\caption{Policy Gradient with MC (REINFORCE)}
	\begin{algorithmic}
    	\STATE initial policy parameters $\theta_0$.
        \FOR {$k=0,1,\dots$}
            \STATE Sample $M$ trajectories $\{\tau_i\}$ according to policy $\pi_{\theta_k}$
            \STATE estimate the gradient via MC: 
            $$
            \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)
            $$
            \STATE update $\theta_k$ using $\hat{g}$ with an optimizer
        \ENDFOR
	\end{algorithmic}
	\end{algorithm}

Variance Reduction

$g$ 虽然是 $\nabla_\theta \mathcal{J}(\theta)$ 的一个无偏估计，但是使用 $g$ 作为估计会出现比较大的方差，我们本节主要介绍如何减少 $g$ 作为估计的方差，这里需要用到 policy gradient baseline invariance 的性质

Baseline Invariance

我们有如下等式

\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]

其中 $b$ 是一个仅与状态有关的函数或者常数。上述等式说明我们可以通过改变 $b$ 来调整 policy gradient 的 variance, 为了证明这个等式，我们证明

\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

注意到

\begin{align} \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{align}

因此，我们就得到了如上结论。 $\blacksquare$

我们将 $b(s_t)$ 称为 baseline.

Reward to Go

Reward to go 的基本思想是，对于时刻 $t$ , 过去的轨迹 $s_0,r_0,\dots,r_{t-1}$ 不应该被考虑，因为它们都已经固定了,加入它们会产生额外的 variance, 此时，我们对应的 $b(s_t)$ 定义为

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

此时，我们的 policy gradient 为

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

对应的 REINFORCE 算法改进版本为

	\begin{algorithm}
	\caption{Policy Gradient with MC and reward to go (REINFORCE)}
	\begin{algorithmic}
    	\STATE initial policy parameters $\theta_0$.
        \FOR {$k=0,1,\dots$}
            \STATE Sample $M$ trajectories $\{\tau_i\}$ according to policy $\pi_{\theta_k}$
            \STATE estimate the gradient via MC: 
            $$
            \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)
            $$
            \STATE update $\theta_k$ using $\hat{g}$ with an optimizer
        \ENDFOR
	\end{algorithmic}
	\end{algorithm}

Optimal Baseline

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

Lemma

令 $s$ 和 $a$ 为随机变量， $w(s,a)$ , $Q(s,a)$ , $b(s)$ 为函数，则

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

其中

b^*(s) = \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]}

证明

\begin{align} \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{align}

其中第三个等式用到了

\begin{align} \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{align}

证毕。 $\blacksquare$

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

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_t)$ 作为 baseline,但是实际上有趣其表达式太复杂，我们一般使用其简化版本，简化版本移除了 $b^*(s)$ 中的 $\nabla_\theta\log\pi_\theta(a_t\mid s_t)$ , 这样我们的一个估计就是

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

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

我们将 $b(s)=V^{\pi_\theta}(s)$ 带入就得到经典的 policy gradient 表达式

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

我们定义

A^{\pi_\theta}(s_t,a_t) := Q^{\pi_\theta}(s_t,a_t)-V^{\pi_\theta}(s_t)

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

如果 $A^{\pi_\theta}(s_t,a_t) >0$ , 则说明 $a_t$ （平均意义下）相对于 $\pi_\theta$ 来说是一个更好的 action
反之如果 $A^{\pi_\theta}(s_t,a_t) <0$ , 则说明 $a_t$ （平均意义下）相对于 $\pi_\theta$ 来说是一个更坏的 action

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

\hat{A}^{\pi_\theta}(s_t,a_t) := Q^{\pi_\theta}(s_t,a_t)-V^{\phi}(s_t)

来代替。

对于最优策略 $\pi^*$ ，我们有

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

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

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

	\begin{algorithm}
	\caption{Policy Gradient with MC and value function (REINFORCE)}
	\begin{algorithmic}
    	\STATE initial policy parameters $\theta_0$.
        \FOR {$k=0,1,\dots$}
            \STATE Sample $M$ trajectories $\{\tau_i\}$ according to policy $\pi_{\theta_k}$
            \STATE estimate the gradient via MC: 
            $$
            \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))
            $$
            \STATE update $\theta_k$ using $\hat{g}$ with an optimizer
        \ENDFOR
	\end{algorithmic}
	\end{algorithm}

Rao–Blackwell Theorem

我们在前面介绍了不同的降低 variance 的方法，但是我们还没有在理论上回答为什么可以降低 variance, 本节我们将基于 Rao-Blackwell theorem 来回答这个问题

Theorem

令 $X$ 和 $Y$ 为随机变量，令 $\hat{I}_1(X,Y)$ 为 $I$ 的一个无偏估计，即

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

令 $\hat{I}_2(Y)=\mathbb{E}_{X\mid Y}[\hat{I}_1(X,Y)\mid Y]$ , 则 $\hat{I}_2$ 也是 $I$ 的一个无偏估计，并且

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

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

证明

由 Expectation 中的 law of total expectation, 我们有

\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

其次，我们有

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

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

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

Overview

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

\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]

则

Method	Expression
Vallinna	$\Psi_t(\tau)=G_0(\tau)=Q^\pi(s_t,a_t)$
Return to go	$\Psi_t(\tau)=\sum_{k=t}^T\gamma^{k-t}r_k$
Advantage	$\Psi_t(\tau)=A^{\pi}(s_t,a_t)=Q^\pi(s_t,a_t)-V^\pi(s_t)$

Experiments

在本节中，我们将对比 REINFORCE 算法的不同变体。我们使用

Conclusion

在本节中，我们推导了 policy gradient 的具体形式，然后针对 REINFORCE 算法 variance 过高的问题，我们介绍了不同的减少 variance 的方法。最后，我们从理论上分析了为什么这些方法可以降低 variance.

Actor Critic Methods

在上一节中，我们介绍了经典的 policy gradient 形式，并基于 policy gradient 构建了经典的 REINFORCE 算法，我们还使用了 baseline invariance 性质来降低 policy gradient 的 variance, 并得到了最终的 advantage 形式。

在本节中，我们将要针对 advantage 进行进一步探讨，介绍如何将 value based methods 和 policy based methods 结合起来得到 actor-critic methods.

Actor-Critic

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

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

因此，其对应的 policy gradient methods 为

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

在前面的章节中，我们介绍了基于 MC 和 TD 两种方式来估计 $Q^{\pi_\theta}(s_t,a_t)$ , 在本节中我们江浙两类方法进行分类：

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

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

	\begin{algorithm}
	\caption{Q actor-critic}
	\begin{algorithmic}
    	\STATE initial policy parameters $\theta_0$, value function parameters $\phi_0$.
        \FOR {$t=0,1,T-1$}
            \STATE $a_t\sim\pi_{\theta}(\cdot\mid s_t)$, $r_t,s_{t+1}\sim p(\cdot,\cdot\mid s_t,a_t)$, $a_{t+1}\sim\pi_{\theta}(\cdot\mid s_{t+1})$ 
            \STATE (ACTOR) policy update
            $$
            \theta_{t+1} = \theta_t + \alpha_{\theta}\log\pi_{\theta}(a_t\mid s_t)Q^{\phi_t}(s_t,a_t)
            $$
            \STATE (CRITIC) value update
            $$
            \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)
            $$
        \ENDFOR
	\end{algorithmic}
	\end{algorithm}

Advantage Actor-Critic (A2C)

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

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

注意到

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

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

\hat{A}_t= r_t + \gamma V_{\phi_t}(s_{t+1})-V_{\phi_t}(s_t)

这样我们就得到了 A2C 算法

	\begin{algorithm}
	\caption{Advantage actor-critic}
	\begin{algorithmic}
    	\STATE initial policy parameters $\theta_0$, value function parameters $\phi_0$.
        \FOR {$t=0,1,T-1$}
            \STATE $a_t\sim\pi_{\theta}(\cdot\mid s_t)$, $r_t,s_{t+1}\sim p(\cdot,\cdot\mid s_t,a_t)$, $a_{t+1}\sim\pi_{\theta}(\cdot\mid s_{t+1})$
            \STATE advantage estimation
            $$
            \hat{A}_t\approx r_t + \gamma V_{\phi_t}(s_{t+1})-V_{\phi_t}(s_t)
            $$ 
            \STATE (ACTOR) policy update
            $$
            \theta_{t+1} = \theta_t + \alpha_{\theta}\log\pi_{\theta}(a_t\mid s_t)\hat{A}_t
            $$
            \STATE (CRITIC) value update
            $$
            \phi_{t+1} = \phi_t + \alpha_\phi\hat{A}_t\nabla_\phi V^{\phi_t}(s_t)
            $$
        \ENDFOR
	\end{algorithmic}
	\end{algorithm}

GAE

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

对比结果如下表所示

Method	bias	variance
TD error	can be biased	low variance
MC estimate	unbiased	high variance

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

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

其中

\delta_t = r_t + \gamma V_{\phi_t}(s_{t+1})-V_{\phi_t}(s_t)

我们可以探究不同的 $\lambda$ 对应的估计：

当 $\lambda=0$ 时，我们有

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

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

当 $\lambda=1$ 时，我们有

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

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

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

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

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

	\begin{algorithm}
	\caption{Generalized Advantage Estimation (GAE)}
	\begin{algorithmic}
    	\STATE policy parameters $\theta$, value function parameters $\phi$.
    	\STATE sample a trajectory $\tau\sim(p_0,\pi_\theta, p)$
        \FOR {$t=0,1,\dots,T-1$}
            \STATE compute TD errors
            $$
            \delta_t = r_t + \gamma V_{\phi}(s_{t+1})-V_{\phi}(s_t)
            $$
        \ENDFOR
        \STATE initialization: $\hat{A}_{T}^{GAE}=0$
        \FOR {$t=T-1,\dots, 0$}
            \STATE iterate backward
            $$
                \hat{A}_t^{GAE} = \delta_t + \gamma\lambda \hat{A}_{t+1}^{GAE}
            $$
        \ENDFOR
	\end{algorithmic}
	\end{algorithm}

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

TRPO

在本文中，我们探讨如何提高 policy gradient methods 的 sample efficiency.

在上一节中介绍的 REINFORCE 算法中，我们每次都需要基于当前策略 $\pi_\theta$ 来进行采样。因此，算法的效率受采样效率所影响。

一个自然的想法是，我们能否利用已有的数据集来进行训练，也就是 off-policy 版本的 policy gradient algorithm.

我们的目标函数为

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

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

\begin{align} \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{align}

我们进一步展开得到

\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}

因此，我们有

\mathcal{J}(\theta) = \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]+C

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

\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}

我们可以证明

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

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

\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(\cdot\mid s)\mid\mid \pi_{\theta_{old}}(\cdot\mid s))\leq \delta \end{align}

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

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

\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

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

Gamma Trick

Overview

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

	\begin{algorithm}
	\caption{TRPO MC}
	\begin{algorithmic}
    	\STATE initial policy parameters $\theta_0$, initial value function parameters $\phi_0$.
    	\STATE hyperparameters: $\delta>0$
        \FOR {$k=0,1,\dots$}
            \STATE Sample $N$ trajectories $\{\tau_i\}$ according to $(p_0,\pi_{\theta_k}, p)$
            \STATE Solve
            $$
            \begin{aligned}
            \max_{\theta_{k+1}}\ &\frac{1}{N}\sum_{i=1}^N\sum_{t=0}^{T^{(i)}-1}\gamma^t\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.}\ &\max_{s\in\mathcal{S}}\mathrm{KL}(\pi_{\theta_{k+1}}(\cdot\mid s)\mid\mid \pi_{\theta_{k}}(\cdot\mid s))\leq \delta
            \end{aligned}
            $$
            \STATE $\hat{R}_t^{(i)}=\gamma^{t'-t}r_{t'}^{(i)}$, $i=1,\dots,N$, $t=1,\dots,T^{(i)}-1$
            \STATE update the value model
            $$
            \min_{\phi} \frac1N\sum_{i=1}^N\frac{1}{T^{(i)}}\sum_{t=0}^{T^{(i)}-1}\frac12\left(V_{\phi}(s_t^{(i)})-\hat{R}_t^{(i)}\right)^2
            $$

        \ENDFOR
	\end{algorithmic}
	\end{algorithm}

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

Experiments

Discussion

Connection with TRO

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

\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}

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

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

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

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

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

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

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

Drawbacks

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

PPO

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

在 TRPO 算法中，我们的优化问题形式为：

\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(\cdot\mid s)\mid\mid \pi_{\theta_{old}}(\cdot\mid s))\leq \delta \end{align}

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

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

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

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

\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]

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

\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}

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

$\hat{A}_t$	$r_t$	objective	gradient	description
$>0$	$r>1+\epsilon$	$(1+\epsilon)\hat{A}_t$	0	好动作，比 old policy 概率更大，不更新
$>0$	$[1-\epsilon,1+\epsilon]$	$r_t\hat{A}_t$	$\nabla_\theta r_t\hat{A}_t$	好动作，与 old policy 概率差不多，更新
$>0$	$r<1-\epsilon$	$(1-\epsilon)\hat{A}_t$	0	好动作，比 old policy 概率更小，不更新
$<0$	$r>1+\epsilon$	$(1+\epsilon)\hat{A}_t$	0	坏动作，比 old policy 概率更大，不更新
$<0$	$[1-\epsilon,1+\epsilon]$	$r_t\hat{A}_t$	$\nabla_\theta r_t\hat{A}_t$	坏动作，比 old policy 概率差不多，更新
$<0$	$r<1-\epsilon$	$(1-\epsilon)\hat{A}_t$	0	坏动作，比 old policy 概率差更小，不更新

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

\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]

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

$\hat{A}_t$	$r_t$	objective	gradient	description
$>0$	$>1+\epsilon$	$(1+\epsilon)\hat{A}_t$	$0$	好动作，比 old policy 概率更大，不需要更新
$>0$	$\leq1+\epsilon$	$r_t\hat{A}_t$	$\nabla_\theta r_t\hat{A}_t$	好动作，与 old policy 概率差不多或更小，需要更新
$<0$	$\geq1-\epsilon$	$r_t\hat{A}_t$	$\nabla_\theta r_t\hat{A}_t$	坏动作，比 old policy 概率更小，不需要概率
$<0$	$<1-\epsilon$	$(1-\epsilon)\hat{A}_t$	$0$	坏动作，与 old policy 概率差不多或更高，概率

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

PPO in RLHF

其中

r_t(\theta) = \frac{\pi_{\theta}(o_t\mid q, o_{< t})}{\pi_{\theta_{old}}(o_t\mid q, o_{< t})}

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

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

其中

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

GRPO

相比于 PPO，GRPO 不依赖于 value function, 因此不需要使用 reward model. GRPO 通过一组输出来估计 value $V(s)$ , 然后进行更新。具体来说，给定 QA pair $(q,a)$ , 我们从 $\pi_{\theta_{old}}$ 中采样 $G$ 个输出 $\{o_i\}_{i=1}^G$ , 接下来我们基于 reward $\{R_i\}_{i=1}^G$ 使用如下表达式来估计 group-level reward:

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

最后，GRPO 的训练目标与 PPO 类似，只不过将 $\hat{A}_t$ 替换为 $\hat{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]

其中，

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})}

KL divergence 的近似见 KL divergence