RL 基础 | Policy Gradient 的推导

去听了 hzxu 老师的 DRL 课，感觉终于听懂了，记录一下…

目录

• 0 我们想做什么
• 1 三个数学 trick
• 2 对单个 transition 的 policy gradient
• 3 对整个 trajectory 的 policy gradient
• 4 REINFORCE 算法

相关链接：

• RL 基础 | Value Iteration 的收敛性证明
• RL 基础 | Policy Iteration 的收敛性证明

---

0 我们想做什么

我们想最大化的东西： $J(\theta )={\mathbb{E}}_{\tau }[R(\tau )]$ ，其中 R 是轨迹的 reward 求和（或 discount 求和）。

我们希望，期望下的轨迹的 reward 求和（reward discounted 求和）最大。

1 三个数学 trick

①： ${\mathrm{\nabla }}_{\theta }\log z=\frac{1}{z}{\mathrm{\nabla }}_{\theta }z$

②： ${\mathbb{E}}_{x\sim p(x)}[f(x)]=\int p(x)f(x)dx$

③： $a/b=[a\cdot p(x)]/[b\cdot p(x)]$

2 对单个 transition 的 policy gradient

$$\begin{array}{rl}{\mathrm{\nabla }}_{\theta }{\mathbb{E}}_{a\sim p(a|s;\theta )}[r(a)]& ={\mathrm{\nabla }}_{\theta }\sum _{a}p(a\mid s;\theta )r(a)\\ & =\sum _{a}r(a){\mathrm{\nabla }}_{\theta }p(a\mid s;\theta )\\ & =\sum _{a}r(a)p(a\mid s;\theta )\frac{{\mathrm{\nabla }}_{\theta }p(a\mid s;\theta )}{p(a\mid s;\theta )}\\ & =\sum _a^{a}r(a)p(a\mid s;\theta ){\mathrm{\nabla }}_{\theta }\log p(a\mid s;\theta )\\ & ={\mathbb{E}}_{a\sim p(a|s;\theta )}[r(a){\mathrm{\nabla }}_{\theta }\log p(a\mid s;\theta )]\end{array}$$

其中，
第一行 把单个 (s,a) 的 reward 期望写为 Σπ(a|s)r(s,a) 的形式；
第二行 认为 r(a) 是不可微分的，去微分 π(a|s)；
第三行 在分数线上下 同时塞了一个 π(a|s) （即 p(a|s;θ) ）；
第四行 因为 d log z = dz/z，原式变成 p(a|s)$\mathrm{\nabla }$p(a|s) 了；
第五行 把 p(a|s) 塞回去，变成了 期望下的 r(s,a) $\mathrm{\nabla }$log π(a|s)。

结论：如果想最大化期望下的 r(s,a)，可以把 r(s,a) 放 $\mathrm{\nabla }$ 外面，去对 log π(a|s) 求梯度。

3 对整个 trajectory 的 policy gradient

先计算 trajectory 的概率：

$$p(\tau \mid \theta )=\underset{\text{initial state distribution}}{\underbrace{\mu (s_0)}}\cdot \prod _{t=0}^{T-1}[\underset{\text{policy}}{\underbrace{\pi (a_t\mid s_t,\theta )}}\cdot \underset{\text{transition fn.}}{\underbrace{p(s_{t+1},r_t\mid s_t,a_t)}}]$$

然后，对单个 transition，我们有

$${\mathrm{\nabla }}_{\theta }{\mathbb{E}}_{x\sim p(x|s;\theta )}[r(x)]={\mathbb{E}}_{x\sim p(x|s;\theta )}[r(x){\mathrm{\nabla }}_{\theta }\log p(x\mid s;\theta )]$$

对于整个 trajectory 的 total reward 的梯度，应用跟 2 相同的方法（分数线上下同乘 p(τ|theta) ），可以得到

$${\mathrm{\nabla }}_{\theta }{\mathbb{E}}_{\tau }[R(\tau )]={\mathbb{E}}_{\tau }[\underset{\text{What is this?}}{\underbrace{{\mathrm{\nabla }}_{\theta }\log p(\tau \mid \theta )}}\underset{\text{Reward of a trajectory}}{\underbrace{R(\tau )}}]$$

现在，让我们来看 ${\mathrm{\nabla }}_{\theta }\log p(\tau \mid \theta )$ 。

$$\begin{array}{rl}\log p(\tau \mid \theta )& =\log \mu (s_0)+\log \prod _{t=0}^{T-1}[\pi (a_t\mid s_t,\theta )\cdot p(s_{t+1},r_t\mid s_t,a_t)]\\ & =\log \mu (s_0)+\sum _{t=0}^{T-1}\log [\pi (a_t\mid s_t,\theta )\cdot p(s_{t+1},r_t\mid s_t,a_t)]\\ & =\log \mu (s_0)+\sum _{t=0}^{T-1}[\log \pi (a_t\mid s_t,\theta )+\log p(s_{t+1},r_t\mid s_t,a_t)]\end{array}$$

其中，
第一行 是把 trajectory 的概率展开；
第二行 第三行 都是把 log(A×B) 变成 logA + logB；
然后发现，只有中间这一项 $\sum _{t=0}^{T-1}\log \pi (a_t\mid s_t,\theta )$ 带 θ，因此，前后两项都不用跟 θ 求梯度了。

由此，我们得到：

$${\mathrm{\nabla }}_{\theta }{\mathbb{E}}_{\tau }[R(\tau )]={\mathbb{E}}_{\tau }[R(\tau ){\mathrm{\nabla }}_{\theta }\sum _{t=0}^{T-1}\log \pi (a_t\mid s_t,\theta )]$$

结论：如果想最大化期望下的 R(τ)，可以把 R(τ) 放 $\mathrm{\nabla }$ 外面，去求 Σ $\mathrm{\nabla }$ log π(a|s) ，即 log [action 概率] 的梯度。

4 REINFORCE 算法

• 使用策略 π(a|s;θ)，生成一个 trajectory：$(s_0,a_0,r_1,...,s_{T-1},a_{T-1},r_T)$ ；
• 对每个时间步 t，计算回报：$R_t=\sum _{k=t+1}^T{\gamma }^{k-t-1}{r}_k$
• 更新策略参数：$\theta =\theta +\alpha {\gamma }^t{R}_t{\nabla }_{\theta }log\pi (a_t|s_t;\theta )$

（算法是 GPT 生成的，看起来好像没问题）