Probability & Matrix

三条公式

$$\begin{array}{rl}{v}_{\pi }& =r_{\pi }+\gamma P_{\pi }{v}_{\pi }\\ r_{\pi }& =\pi @P_{r}R\\ P_{\pi }& =\pi @P_s\end{array}$$

将概率论和线性代数统一起来

对于下面的Bellman Equation，实际上每个概率我们都可以写成matrix form.

$$\begin{array}{rl}{v}_{\pi }(s)& =r_{\pi }(s)+\gamma \sum _{s^{\prime }}p(s^{\prime }|s)v_{\pi }(s^{\prime })\\ r_{\pi }(s)& \triangleq \sum _a\pi (a|s)\sum _{r}p(r|s,a)r\\ p_{\pi }(s^{{}^{\prime }}|s)& \triangleq \sum _a\pi (a|s)p(s^{{}^{\prime }}|s,a)\end{array}$$

我们分两种情况考虑，一种是普通的probability，另外一种则是条件概率 conditional probability，对于前者：

$$v_{\pi }=r_{\pi }+\gamma P_{\pi }{v}_{\pi }$$

where

• $v_{\pi }=[v_{\pi }(s_1),v_{\pi }(s_2),\dots ,v_{\pi }(s_n){]}^T\in {\mathbb{R}}^n$
• $r_{\pi }=[r_{\pi }(s_1),r_{\pi }(s_2),\dots ,r_{\pi }(s_n){]}^T\in {\mathbb{R}}^n$
• $P_{\pi }\in {\mathbb{R}}^{n\times n}$, where $[P_{\pi }{]}_{ij}=p_{\pi }(s_j\mid s_i)={\sum }_a\pi (a\mid s_i)p(s_j\mid s_i,a)$

对于前两点，我们很容易通过堆叠多个$s$的方法从而向量化：

• $v_{\pi }(s)\to v_{\pi }\in {\mathbb{R}}^n$
• $r_{\pi }(s)\to r_{\pi }\in {\mathbb{R}}^n$

但是却没有向量化彻底，对于三个分布我们都还没有向量化，而刚好这三个分布又都是条件概率的形式，对于这种特殊的格式，而且我们往往有这样的潜规则：

$$p(\text{output}\mid \text{input})\to \text{input\_size}\{\begin{array}{c}\\ \\ \\ \end{array}\underset{\text{output\_size}}{\underbrace{\left[\begin{array}{ccc}\cdot & \cdot & \cdot \\ \cdot & \cdot & \cdot \\ \cdot & \cdot & \cdot \\ \cdot & \cdot & \cdot \end{array}\right]}}$$

where $\text{input}\in {\mathbb{R}}^n$ and $\text{output}\in {\mathbb{R}}^m$, and we have ${\sum }_{\text{output}}p(\text{output}\mid \text{input})=1$, so that the matrix P on the right side must have P.sum(dim=1)=[1,1,...,1].T

那么很简单的我们就有如下几个向量化之后的分布

• $\pi (a\mid s)\to \pi \in {\mathbb{R}}^{n\times m}$
  • where we code it policy, policy.shape=(n,m)
• $p(s^{\prime }\mid s,a)\to P_s\in {\mathbb{R}}^{n\times m\times n}$
  • where we code it Ps, Ps.shape=(n,m,n), Ps.sum(dim=-1)=[[1,1,...],[1,1,...],...]
• $p(r\mid s,a)\to P_r\in {\mathbb{R}}^{n\times m\times k}$
  • should recall that $R\in {\mathbb{R}}^k$, means the set of reward but we vectorize it

那么我们接下来把剩下两条公式也写成matrix-form

$$\begin{array}{rl}{r}_{\pi }& =\pi @P_{r}R\\ P_{\pi }& =\pi @P_s\end{array}$$

where

• $@$ 代表了batch matrix multiply (bmm)
• $\pi @P_s$实际上是policy先扩展一个维度变成(n,1,m) @ (n,m,n)，这样得出的结果是(n,1,n)，再squeeze就得到了(n,n)
  • 另外由于policy.sum(-1)=[1,1,...,1].T, Ps.sum(-1)=[[1,1,...],...]，最后得到的结果P_pi其实也是满足 P_pi..sum(-1)=[1,1,...,1](简单展开可证)
• $\pi @P_{r}R$的操作是(n,1,m) @ (n,m,k) (k,) \
  • 我们关注点就是原公式的求和符号的下标，下标再执行求和符号之后是会消掉的
  • 矩阵相乘也是一样：
    • (n,m,k)(k,) = (n,m)其实就是k消掉了，此时执行了${\sum }_r$
    • (n,1,m)@(n,m,1)=(n,1)其实就是m消掉了，此时执行了${\sum }_a$

一些联想

• $\pi \in {\mathbb{R}}^{n\times m}$ 正好就是我们之后要学习的Q table或者是网络的参数 （我们假设action space大小为m）
• 这里的$P_r,P_s$代表着环境model，都是环境给的；但这里的下标是为了我们方便看出来是和state有关还是和reward有关，之前的出现的$P_{\pi },r_{\pi }$的下标其实是base的意思，也就是说$r_{\pi }(s)=r(s,\pi )$