optimal_baseline

Derivation of optimal baseline

Preliminiary

1. Trace of Matrix
2. here we have $\mathbb{E}={\mathbb{E}}_{S\sim \mu ,A\sim \pi }$

The gradient is ${\nabla }_{\theta }J(\theta )=\mathbb{E}(X)$ where

$$X\doteq (q_{\pi }(S,A)-b(S)){\nabla }_{\theta }\ln \pi (A|S,\theta )$$

• $X$ is a (k,1) vector, same as the gradient ${\nabla }_{\theta }\ln \pi (A|S,\theta )$
• note ${\nabla }_{\theta }J(\theta )$ or $\theta$ is also a (k,1) vector
• Let $\bar{x}\doteq \mathbb{E}(X)$, also a (k,1) vector

It is common to use trace to measure the total variance of a covariance matrix $\mathrm{var}(X)$, and trace is also a scalar obejective function to be optimized:

$$\begin{array}{rlrl}\mathrm{tr}\left[\mathrm{var}(X)\right]& =\mathrm{tr}\left[\mathbb{E}\left[(X-\bar{x})(X-\bar{x}{)}^T\right]\right]& & \\ & =\mathrm{tr}\left[\mathbb{E}\left[\underset{(k,k)}{\underbrace{X{X}^T}}-\underset{(k,k)}{\underbrace{\bar{x}{X}^T}}-\underset{(k,k)}{\underbrace{X{\bar{x}}^T}}+\underset{(k,k)}{\underbrace{\bar{x}{\bar{x}}^T}}\right]\right]& & \\ & =\mathbb{E}\left[\underset{(1,1)}{\underbrace{X^{T}X}}-\underset{(1,1)}{\underbrace{X^T\bar{x}}}-\underset{(1,1)}{\underbrace{{\bar{x}}^{T}X}}+\underset{(1,1)}{\underbrace{{\bar{x}}^T\bar{x}}}\right]& & (\mathrm{tr}(AB)=\mathrm{tr}(BA))\\ & =\mathbb{E}\left[X^{T}X\right]-{\bar{x}}^T\bar{x}& & \end{array}$$

Then we have

$$\begin{array}{rl}{\nabla }_b\mathrm{tr}\left[\mathrm{var}(X)\right]& ={\nabla }_b\left(\mathbb{E}\left[X^{T}X\right]-{\bar{x}}^T\bar{x}\right)\\ & ={\nabla }_b\mathbb{E}\left[X^{T}X\right]-\underset{0}{\underbrace{{\nabla }_b{\bar{x}}^T\bar{x}}}\\ & ={\nabla }_b\mathbb{E}\left[\underset{(1,k)}{\underbrace{({\nabla }_{\theta }\ln \pi {)}^T}}\underset{(k,1)}{\underbrace{({\nabla }_{\theta }\ln \pi )}}(q_{\pi }(S,A)-b(S){)}^2\right]\\ & ={\nabla }_b\mathbb{E}\left[\Vert {\nabla }_{\theta }\ln \pi {\Vert }^2(q_{\pi }(S,A)-b(S){)}^2\right]\\ & =-2\mathbb{E}\left[\Vert {\nabla }_{\theta }\ln \pi {\Vert }^2(q_{\pi }(S,A)-b(S))\right]\end{array}$$

To optimize the variance, we need to set the gradient to zero:

$$\begin{array}{rl}{\nabla }_b\mathrm{tr}\left[\mathrm{var}(X)\right]& =\mathbb{E}\left[\Vert {\nabla }_{\theta }\ln \pi {\Vert }^2(q_{\pi }(S,A)-b(S))\right]=0\end{array}$$

To ensure the above equation holds, $b(s)$ for any $s\in \mathcal{S}$ should satisfy:

$${\mathbb{E}}_{A\sim \pi }\left[\Vert {\nabla }_{\theta }\ln \pi (A|s,\theta ){\Vert }^2(q_{\pi }(s,A)-b(s))\right]=0$$

• note: we just expand the expectation $\mathbb{E}$ to state-wise form

The above equation can be easily solved to obtain the optimal baseline:

$$\boxed{b^{\ast }(s)=\frac{{\mathbb{E}}_{A\sim \pi }\left[q_{\pi }(s,A)\Vert {\nabla }_{\theta }\ln \pi (A|s,\theta ){\Vert }^2\right]}{{\mathbb{E}}_{A\sim \pi }\left[\Vert {\nabla }_{\theta }\ln \pi (A|s,\theta ){\Vert }^2\right]}}$$

We can remove the weights $\Vert {\nabla }_{\theta }\ln \pi (A|s,\theta ){\Vert }^2$ from the above equation, and obtain the suboptimal baseline in a more concise form:

$$\boxed{b^{\ast }(s)={\mathbb{E}}_{A\sim \pi }[q_{\pi }(s,A)]=v_{\pi }(s)}$$

说一些证明过程中的感悟

这一章的证明其实不难，但是自己依然花了比较多的时间，其原因是自己一开始就把X的shape理解错成了(n,m)的matrix，原因很简单，就是自认为和之前在policy_gradient_proof中的$\nabla \pi$的思想一样，思路是$\pi$.shape=(n,m)⇒$\nabla \pi$.shape=(n,m)⇒$\nabla \ln \pi$.shape=(n,m)，但实际上直接这样想是错误的，原因如下：

• 在TL;DR的章节中的$\nabla$是基本按$d$的操作的格式去一路写下去的，这也是为什么自己可以这样写：$\nabla v_{\pi }=\nabla (\pi @q_{\pi })=\nabla \pi @q_{\pi }+\pi @\nabla q_{\pi }$，这个过程是符合$\pi$.shape=(n,m)⇒$\nabla \pi$.shape=(n,m)的，这里的$\nabla \pi =d\pi$，但是policy_gradient_proof中除了TLDR这里的推导($\nabla \ne {\nabla }_{\theta }$)以外，其他所有地方的$\nabla \pi$.shape都是等于$\theta$.shape的（因为最开头明确说明了$\nabla ={\nabla }_{\theta }$）
• 本章节目的是针对参数$\theta$(e.g. (k,1) vector)，希望可以找到一条梯度下降的方法，既不改变原本下降的梯度${\nabla }_{\theta }J(\theta )=\mathbb{E}(X)$，也可以使得对于$\theta$而言这条路径梯度的方差$\mathrm{var}$不会太大（上面我们就是用trace将(k,k)的这个协方差矩阵的对角线上的协方差值加起来，正好变成一个scalar，去进而求关于b的梯度，得到b的最优解格式），这里的梯度应该是符合$\theta$.shape=(k,1)⇒${\nabla }_{\theta }\pi$.shape=(k,1)

总结就是说，X.shape=theta.shape=(k,1)才是正解，最开头关于X的定义式是element-wise的（即对于每个state每个action而言我们是展开写的式子，这里并不可以一眼就把$q_{\pi }(S,A)$直接就写成(n,m)矩阵$q_{\pi }$了，这样做下去会把问题复杂化）可以回顾这一节内容⇒ About the policy gradient