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

①： \(\nabla_\theta\log z = \frac1z\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{aligned} \nabla_\theta\mathbb{E}_{a\sim p(a|s;\theta)}[r(a)]& =\nabla_\theta\sum_ap(a\mid s;\theta)r(a) \\ &=\sum_ar(a)\nabla_\theta p(a\mid s;\theta) \\ &=\sum_ar(a)p(a\mid s;\theta)\frac{\nabla_\theta p(a\mid s;\theta)}{p(a\mid s;\theta)} \\ &=\sum_a^ar(a)p(a\mid s;\theta)\nabla_\theta\log p(a\mid s;\theta) \\ &=\mathbb{E}_{a\sim p(a|s;\theta)}[r(a)\nabla_\theta\log p(a\mid s;\theta)] \end{aligned} \]

其中，
第一行 把單個 (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)\(\nabla\)p(a|s) 了；
第五行 把 p(a|s) 塞回去，變成了 期望下的 r(s,a) \(\nabla\)log π(a|s)。

結論：如果想最大化期望下的 r(s,a)，可以把 r(s,a) 放 \(\nabla\) 外面，去對 log π(a|s) 求梯度。

3 對整個 trajectory 的 policy gradient

先計算 trajectory 的機率：

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

然後，對單個 transition，我們有

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

對於整個 trajectory 的 total reward 的梯度，應用跟 2 相同的方法（分數線上下同乘 p(τ|theta) ），可以得到

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

現在，讓我們來看 \(\nabla_\theta\log p(\tau\mid\theta)\) 。

\[\begin{aligned} \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{aligned} \]

其中，
第一行 是把 trajectory 的機率展開；
第二行 第三行 都是把 log(A×B) 變成 logA + logB；
然後發現，只有中間這一項 \(\sum_{t=0}^{T-1}\log\pi(a_t\mid s_t,\theta)\) 帶 θ，因此，前後兩項都不用跟 θ 求梯度了。

由此，我們得到：

\[\nabla_\theta\mathbb{E}_\tau[R(\tau)]=\mathbb{E}_\tau\left[R(\tau)\nabla_\theta\sum_{t=0}^{T-1}\log\pi(a_t\mid s_t,\theta)\right] \]

結論：如果想最大化期望下的 R(τ)，可以把 R(τ) 放 \(\nabla\) 外面，去求 Σ \(\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} γ^{k-t-1} r_k\)
• 更新策略引數：\(θ = θ + α γ^t R_t ∇_θ log π(a_t|s_t;θ)\)

（演算法是 GPT 生成的，看起來好像沒問題）