基础的 TD 和 RL 概念已在 RL-basis 中介绍过,这里不再重复。

n-step return

在TD(0) 里面有: $\delta_t := R_{t+1} + \gamma V(S_{t + 1}) - V(S_t)$ 可以看作是单步奖励(确定) + 对下一步的期望 - 当前状态的期望,这是单步的TD error
首先引出n-step return:
单步return 就是很简单的 $G_t = R_{t + 1} + \gamma V(S_{t+1})$
很自然有n-step return是: $G_t^{(n)} = R_{t + 1} + \gamma R_{t + 2} + \gamma^2 R_{t + 3}+ \dots +\gamma^{n - 1} R_{t + n} + \gamma^n V(S_{t + n})$.
如果我们展开

$$ \begin{aligned} \sum_{k = t}^{t + n - 1} \gamma^{k-t} \delta_k &= (R_{t+1} + \gamma V(S_{t+1}) - V(S_t)) \\\\ &\quad + \gamma(R_{t+2} + \gamma V(S_{t+2}) - V(S_{t+1})) \\\\ &\quad + \dots \\\\ &\quad + \gamma^{n-1}(R_{t+n} + \gamma V(S_{t+n}) - V(S_{t+n-1})) \\\\ &= G_t^{(n)} - V(S_t) \end{aligned} $$

移项即得 $G_t^{(n)} = V(S_t) + \sum_{k=t}^{t+n-1} \gamma^{k-t} \delta_k$:n-step return = 当前状态的估计值 + 未来 $n$ 步 TD error 的修正。

n-step return 的意义在于它是 MC 和 TD 之间的 trade-off——MC 无偏但方差大,TD 有偏但方差小。n 越大越偏向 MC,n 越小越偏向 TD。

$\lambda - Return$

$$ \begin{aligned} G_t^\lambda &= (1 - \lambda) \sum_{n=1}^{\infty} \lambda^{n - 1} G_t^{(n)} && \text{其中 } 1-\lambda \text{ 是归一化系数} \\\\ &= V(S_t) + \sum_{k = 0}^{\infty} ({\gamma \lambda})^k \delta_{t_k} \end{aligned} $$

$$ \begin{cases} \text{TD(0)} &\lambda = 0 \\\\ \text{MC} &\lambda = 1 \\\\ \text{trade-off} &\lambda \in (0,1) \end{cases} $$

上述 n-step return 和 $\lambda$-return 都属于 value-based 方法。这类方法有两个主要局限:

  • 维度爆炸:当 trajectory 过长或 state/action 空间很大时,tabular 方法无法应对,每一步还要做 max 操作
  • 仅适用于离散动作:连续动作空间下无法对每个 action 求 max

而 policy-based 方法直接参数化策略 $\pi_\theta$,输出的是动作的概率分布,避免了 max 操作,天然支持连续空间,同时保留了随机性。

policy gradient

首先我们是为了得到在给定的一个policy里面让我们的 $V_{\pi \sim \theta}$ 尽可能大,因此我们对其进行展开.

$$ \begin{aligned} V^{\pi}(s) &= \sum_{a} \pi_{\theta}(a|s) Q^{\pi}(s,a) && \text{对其两边进行求导运算} \\\\ \nabla V^{\pi}(s) &= \sum_{a} \underbrace{\nabla \pi_{\theta}(a|s) Q^{\pi}(s,a)}_{\text{记为 }\phi(s)} + \pi_{\theta}(a|s) \nabla Q^{\pi}(s,a) \end{aligned} $$

其中 $\phi(s)$ 可以看作当前 step 的 reward 贡献;后面那一项更关键,因为 $Q^{\pi}$ 会通过轨迹对后续所有 $S, A, R$ 产生连锁影响。

用 Bellman 方程展开 $Q^{\pi}$:

$$ \begin{aligned} Q^{\pi}(s,a) &= R(s,a) + \gamma \sum_{s'} P(s' \mid s,a) V^{\pi}(s') \\\\ \nabla Q^{\pi}(s,a) &= \gamma \sum_{s'} P(s' \mid s,a) \nabla V^{\pi}(s') \end{aligned} $$

可以看到 $\nabla V^{\pi}$ 在等式两边都出现了,因此可以递归求解。

引入 1-step visitation probability, 记 $\rho^{\pi}(s \to s', 1) = \sum_{a} \pi_{\theta}(a|s) P(s'|s,a)$, 表示在策略 $\pi$ 下从 $s$ 走一步到达 $s'$ 的概率. 把 $\nabla Q^{\pi}$ 代入 $\nabla V^{\pi}$ 后就可以写成:

$$ \nabla V^{\pi}(s) = \phi(s) + \sum_{s'} \rho^{\pi}(s \to s', 1) \nabla V^{\pi}(s') $$

现在对 $\nabla V^{\pi}$ 不断展开(unroll), 每次把递归的部分再代一次同样的公式:

$$ \begin{aligned} & \nabla V^{\pi}(s) \\ &= \phi(s) + \sum_{a} \pi_{\theta}(a|s) \sum_{s'} P(s'|s,a) \nabla V^{\pi}(s') \\\\ &= \phi(s) + \sum_{s'} \rho^{\pi}(s \to s', 1) \nabla V^{\pi}(s') \\\\ &= \phi(s) + \sum_{s'} \rho^{\pi}(s \to s', 1) \Big[ \phi(s') + \sum_{s''} \rho^{\pi}(s' \to s'', 1) \nabla V^{\pi}(s'') \Big] \\\\ &= \phi(s) + \sum_{s'} \rho^{\pi}(s \to s', 1) \phi(s') + \sum_{s''} \underbrace{\sum_{s'} \rho^{\pi}(s \to s', 1) \rho^{\pi}(s' \to s'', 1)}_{\rho^{\pi}(s \to s'', 2)} \nabla V^{\pi}(s'') \\\\ &= \phi(s) + \sum_{s'} \rho^{\pi}(s \to s', 1) \phi(s') + \sum_{s''} \rho^{\pi}(s \to s'', 2) \phi(s'') + \sum_{s'''} \rho^{\pi}(s \to s''', 3) \nabla V^{\pi}(s''') \\\\ &= \dots \end{aligned} $$

其中 $\rho^{\pi}(s \to x, k)$ 是 k-step visitation probability: 从 $s$ 出发, 在策略 $\pi$ 下恰好走 $k$ 步到达 $x$ 的概率. 它满足递推关系 $\rho^{\pi}(s \to x, k+1) = \sum_{s'} \rho^{\pi}(s \to s', k) \rho^{\pi}(s' \to x, 1)$.

一直展开下去, $\nabla V^{\pi}$ 的递归项不断后移, 最终:

$$ \nabla V^{\pi}(s) = \sum_{x \in \mathcal{S}} \sum_{k=0}^{\infty} \rho^{\pi}(s \to x, k) \phi(x) $$

这就把所有的递归都吃掉了, $\nabla V^{\pi}$ 写成了 $k$-step visitation 的加权和.

接下去把目标函数 $J(\theta)$ 引入. $J(\theta)$ 就是最大化策略的期望回报, 从状态 $s_0$ 出发就等于 $V^{\pi}(s_0)$. 记 $\eta(s) = \sum_{k=0}^{\infty} \rho^{\pi}(s_0 \to s, k)$, 它表示从 $s_0$ 出发, 在所有时间步内访问状态 $s$ 的总概率(未归一化). 于是:

$$ \begin{aligned} \nabla J(\theta) &= \nabla V^{\pi}(s_0) \\\\ &= \sum_{s} \underbrace{\sum_{k=0}^{\infty} \rho^{\pi}(s_0 \to s, k)}_{\eta(s)} \phi(s) \\\\ &= \sum_{s} \eta(s) \phi(s) \\\\ &= \Big( \sum_{s} \eta(s) \Big) \sum_{s} \frac{\eta(s)}{\sum_{s} \eta(s)} \phi(s) \\\\ &\propto \sum_{s} d^{\pi}(s) \phi(s) \end{aligned} $$

其中 $d^{\pi}(s) = \frac{\eta(s)}{\sum_{s} \eta(s)}$ 是归一化后的 stationary distribution(满足 $\sum_s d^{\pi}(s) = 1$),$\sum_s \eta(s)$ 是常数项,在梯度中被 $\propto$ 吸收掉。

$\propto$ 的常数到底是多少?

Episodic 设定(本文的推导路线)

$$ \begin{aligned} \sum_{s} \eta(s) &= \sum_{s} \sum_{k=0}^{\infty} \rho^{\pi}(s_0 \to s, k) = \sum_{k=0}^{\infty} \underbrace{\sum_{s} \rho^{\pi}(s_0 \to s, k)}_{P(T > k)} \\ &= \sum_{k=0}^{\infty} P(T > k) = \mathbb{E}[T] \end{aligned} $$

对每个固定的第 $k$ 步,$\sum_s \rho^{\pi}(s_0 \to s, k)$ 是所有状态上的概率之和,表达的是「episode 在第 $k$ 步仍存活」的概率,不是 1(因为部分概率 mass 已被终止态吸收)。根据离散随机变量的期望公式 $\mathbb{E}[T] = \sum_{k=0}^{\infty} P(T > k)$,有 $\sum_s \eta(s) = \mathbb{E}[T]$,正好是平均 episode 长度。因此:

$$ \nabla J(\theta) = \mathbb{E}[T] \cdot \sum_s d^{\pi}(s) \phi(s) \;\propto\; \sum_s d^{\pi}(s) \phi(s) $$

Continuing (average-reward) 设定

此时没有终止态,$d^{\pi}(s)$ 直接定义为 Markov chain 在策略 $\pi$ 下的平稳分布,它是转移概率矩阵的特征向量,天然满足 $\sum_s d^{\pi}(s) = 1$,不需要手动归一化。Policy gradient theorem 在 average-reward 框架下证明出来就是等式,常数为 1(Sutton & Barto, 2017, Sec. 13.2)。

两种设定最终形式一致,只是常数的含义不同。

把 $\phi(s) = \sum_{a} \nabla \pi_{\theta}(a|s) Q^{\pi}(s,a)$ 代回去:

$$ \nabla J(\theta) \propto \sum_{s \in \mathcal{S}} d^{\pi}(s) \sum_{a \in \mathcal{A}} \nabla \pi_{\theta}(a|s) Q^{\pi}(s,a) $$

score function

最后一步,用 score function(也叫 likelihood ratio)把 $\nabla \pi$ 化掉:

$$ \begin{aligned} \nabla J(\theta) &\propto \sum_{s \in \mathcal{S}} d^{\pi}(s) \sum_{a \in \mathcal{A}} Q^{\pi}(s,a) \nabla \pi_{\theta}(a|s) \\\\ &= \sum_{s \in \mathcal{S}} d^{\pi}(s) \sum_{a \in \mathcal{A}} \pi_{\theta}(a|s) Q^{\pi}(s,a) \frac{\nabla \pi_{\theta}(a|s)}{\pi_{\theta}(a|s)} \\\\ &= \mathbb{E}_{s \sim d^{\pi}, a \sim \pi_{\theta}} \left[ Q^{\pi}(s,a) \nabla \ln \pi_{\theta}(a|s) \right] \qquad \text{; } (\ln x)' = \frac{1}{x} \end{aligned} $$

这就是经典的 Policy Gradient Theorem——期望里面不再有求和。梯度方向由 $Q$ 加权,沿着 $\nabla \ln \pi$ 更新策略参数。实际中 $Q^{\pi}$ 可用 Monte Carlo return $G_t$、TD target、advantage $A^{\pi}$ 等方式估计。

References