Sources:

1. Shiyu Zhao. Chapter 2: State Values and Bellman Equation. Mathematical Foundations of Reinforcement Learning.
2. --> Youtube: Bellman Equation: Matrix-Vector form

Bellman equation: the matrix-vector form

Consider the Bellman equation: $$v_{\pi }(s)=\sum _a\pi (a\mid s)[\sum _{r}p(r\mid s,a)r+\gamma \sum _{s^{\prime }}p(s^{\prime }\mid s,a)v_{\pi }\left(s^{\prime }\right)]$$

It's an elementsise form. That means there are $|\mathcal{S}|$ equations like this! If we put all the equations together, we have a set of linear equations, which can be concisely written in a matrix-vector form.

Recall that this equation can be rewritten as $$v_{\pi }(s)=r_{\pi }(s)+\gamma \sum _{s^{\prime }}{p}_{\pi }(s^{\prime }\mid s)v_{\pi }\left(s^{\prime }\right)$$ where $$r_{\pi }(s)\triangleq \sum _a\pi (a\mid s)\sum _{r}p(r\mid s,a)r,p_{\pi }(s^{\prime }\mid s)\triangleq \sum _a\pi (a\mid s)p(s^{\prime }\mid s,a)$$

Suppose the states could be indexed as $s_i(i=1,\dots ,n)$. For state $s_i$, the Bellman equation is $$v_{\pi }\left(s_i\right)=r_{\pi }\left(s_i\right)+\gamma \sum _{s_j}{p}_{\pi }(s_j\mid s_i)v_{\pi }\left(s_j\right)$$

Put all these equations for all the states together and rewrite to a matrix-vector form $$\begin{array}{}\text{(1)}& v_{\pi }=r_{\pi }+\gamma P_{\pi }{v}_{\pi }\end{array}$$ where

$v_{\pi }={[v_{\pi }\left(s_1\right),\dots ,v_{\pi }\left(s_n\right)]}^T\in {\mathbb{R}}^n$,

$r_{\pi }={[r_{\pi }\left(s_1\right),\dots ,r_{\pi }\left(s_n\right)]}^T\in {\mathbb{R}}^n$,

$P_{\pi }\in {\mathbb{R}}^{n\times n}$, where ${\left[P_{\pi }\right]}_{ij}=p_{\pi }(s_j\mid s_i)$, is the state transition matrix.

Examples

Refer to the grid world example for the notations.

If there are four states, $v_{\pi }=r_{\pi }+\gamma P_{\pi }{v}_{\pi }$ can be written out as $$\underset{v_{\pi }}{\underbrace{\left[\begin{array}{l}{v}_{\pi }\left(s_1\right)\\ v_{\pi }\left(s_2\right)\\ v_{\pi }\left(s_3\right)\\ v_{\pi }\left(s_4\right)\end{array}\right]}}=\underset{r_{\pi }}{\underbrace{\left[\begin{array}{l}{r}_{\pi }\left(s_1\right)\\ r_{\pi }\left(s_2\right)\\ r_{\pi }\left(s_3\right)\\ r_{\pi }\left(s_4\right)\end{array}\right]}}+\gamma \underset{P_{\pi }}{\underbrace{\left[\begin{array}{llll}{p}_{\pi }(s_1\mid s_1)& p_{\pi }(s_2\mid s_1)& p_{\pi }(s_3\mid s_1)& p_{\pi }(s_4\mid s_1)\\ p_{\pi }(s_1\mid s_2)& p_{\pi }(s_2\mid s_2)& p_{\pi }(s_3\mid s_2)& p_{\pi }(s_4\mid s_2)\\ p_{\pi }(s_1\mid s_3)& p_{\pi }(s_2\mid s_3)& p_{\pi }(s_3\mid s_3)& p_{\pi }(s_4\mid s_3)\\ p_{\pi }(s_1\mid s_4)& p_{\pi }(s_2\mid s_4)& p_{\pi }(s_3\mid s_4)& p_{\pi }(s_4\mid s_4)\end{array}\right]}}\underset{v_{\pi }}{\underbrace{\left[\begin{array}{l}{v}_{\pi }\left(s_1\right)\\ v_{\pi }\left(s_2\right)\\ v_{\pi }\left(s_3\right)\\ v_{\pi }\left(s_4\right)\end{array}\right]}}.$$

For deterministic policy

Figure 2.4: An example for demonstrating the Bellman equation. The policy in this example is deter- ministic.

For this specific example: $$\left[\begin{array}{l}{v}_{\pi }\left(s_1\right)\\ v_{\pi }\left(s_2\right)\\ v_{\pi }\left(s_3\right)\\ v_{\pi }\left(s_4\right)\end{array}\right]=\left[\begin{array}{l}0\\ 1\\ 1\\ 1\end{array}\right]+\gamma \left[\begin{array}{llll}0& 0& 1& 0\\ 0& 0& 0& 1\\ 0& 0& 0& 1\\ 0& 0& 0& 1\end{array}\right]\left[\begin{array}{l}{v}_{\pi }\left(s_1\right)\\ v_{\pi }\left(s_2\right)\\ v_{\pi }\left(s_3\right)\\ v_{\pi }\left(s_4\right)\end{array}\right]$$

For stochastic policy

Figure 2.5: An example for demonstrating the Bellman equation. The policy in this example is stochastic.

For this specific example: $$\left[\begin{array}{c}{v}_{\pi }\left(s_1\right)\\ v_{\pi }\left(s_2\right)\\ v_{\pi }\left(s_3\right)\\ v_{\pi }\left(s_4\right)\end{array}\right]=\left[\begin{array}{c}0.5(0)+0.5(-1)\\ 1\\ 1\\ 1\end{array}\right]+\gamma \left[\begin{array}{cccc}0& 0.5& 0.5& 0\\ 0& 0& 0& 1\\ 0& 0& 0& 1\\ 0& 0& 0& 1\end{array}\right]\left[\begin{array}{l}{v}_{\pi }\left(s_1\right)\\ v_{\pi }\left(s_2\right)\\ v_{\pi }\left(s_3\right)\\ v_{\pi }\left(s_4\right)\end{array}\right]$$

Solution of the matrix-vector form

Recalling the Bellman equation in matrix-vector form $\text{(1)}$,

We can convert it to two forms:

• The closed-form solution is: $$v_{\pi }={(I-\gamma P_{\pi })}^{-1}{r}_{\pi }$$ The drawback of closed-form solution is that it involves a matrix inverse operation, which is computationally expensive. Thus, in practice, we'll use an iterative solution.

• An iterative solution is: $$v_{k+1}=r_{\pi }+\gamma P_{\pi }{v}_k,$$ where $I$ is the identity matrix.

  We can just randomly select a matrix $v_0$, then calculate $v_1,v_2,\cdots$. This leads to a sequence $\{v_0,v_1,v_2,\dots \}$. We can show that $$v_k\to v_{\pi }={(I-\gamma P_{\pi })}^{-1}{r}_{\pi },k\to \mathrm{\infty }$$

Proof: the closed-form solution

First, the Bellman equation in matrix-vector form is $$v_{\pi }=r_{\pi }+\gamma P_{\pi }{v}_{\pi }.$$ Put the $\gamma P_{\pi }{v}_{\pi }$ into the left side: $$\begin{array}{r}{v}_{\pi }-\gamma P_{\pi }{v}_{\pi }=r_{\pi }\\ (I-\gamma P_{\pi })v_{\pi }=r_{\pi }\end{array}$$

Now we calculate the matrix inverse: $$v_{\pi }={(I-\gamma P_{\pi })}^{-1}{r}_{\pi }.$$ Q.E.D.

Proof: the iterative solution

Define the error as ${\delta }_k=v_k-v_{\pi }$. We only need to show ${\delta }_k\to 0$. Substituting:

1. $v_{k+1}={\delta }_{k+1}+v_{\pi }$ and
2. $v_k={\delta }_k+v_{\pi }$

into $v_{k+1}=r_{\pi }+\gamma P_{\pi }{v}_k$ gives $${\delta }_{k+1}+v_{\pi }=r_{\pi }+\gamma P_{\pi }({\delta }_k+v_{\pi })$$ which can be rewritten as $${\delta }_{k+1}=-v_{\pi }+r_{\pi }+\gamma P_{\pi }{\delta }_k+\gamma P_{\pi }{v}_{\pi }=\gamma P_{\pi }{\delta }_k$$

As a result, $${\delta }_{k+1}=\gamma P_{\pi }{\delta }_k={\gamma }^2{P}_{\pi }^2{\delta }_{k-1}=\cdots ={\gamma }^{k+1}{P}_{\pi }^{k+1}{\delta }_0$$

Note that $0\le P_{\pi }^k\le 1$, which means every entry of $P_{\pi }^k$ is no greater than 1 for any $k=0,1,2,\dots$. That is because $P_{\pi }^k\mathbf{1}=\mathbf{1}$, where $\mathbf{1}=[1,\dots ,1{]}^T$. On the other hand, since $\gamma <1$, we know ${\gamma }^k\to 0$ and hence $${\delta }_{k+1}={\gamma }^{k+1}{P}_{\pi }^{k+1}{\delta }_0=\gamma P_k^{\pi }({\gamma }^k{P}_{\pi }^k{\delta }_0)\to 0$$ as $k\to \mathrm{\infty }$.