RL 基础 | Policy Gradient 的推导

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

目录

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

相关链接:


0 我们想做什么

我们想最大化的东西: \(J(\theta) = \mathbb E_\tauR(\\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\leftR(\\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 生成的,看起来好像没问题)