Bellman equation derivation

Bellman equation for $v_{\pi }$.

We want to show $v_{\pi }(s)={\sum }_a\pi (a\mid s){\sum }_{s^{\prime },r}p(s^{\prime },r\mid s,a)[r+\gamma v_{\pi }(s^{\prime })]$ for all states $s$.

The core idea of the proof is to use the law of total probability to go from marginal to conditional probabilities, and then invoke the Markov assumption.

The law of total probability states that if $B$ is an event, and $C_1,\dots ,C_n$ are events that partition the sample space, then $Pr(B)=\sum _{j=1}^{n}Pr(B\cap C_j)=\sum _{j=1}^{n}Pr(B\mid C_j)Pr(C_j)$.

For fixed event $A$ with non-zero probability, the mapping $B\mapsto Pr(B\mid A)$ is another valid probability measure. In other words, define ${Pr}_A$ by ${Pr}_A(B):=Pr(B\mid A)$ for all events $B$. Now the law of total probability for $P_A$ states that ${Pr}_A(B)=\sum _{j=1}^n{Pr}_A(B\mid C_j){Pr}_A(C_j)$. We also have

${Pr}_A(B\mid C_j)=\frac{{Pr}_A(B\cap C_j)}{{Pr}_A(C_j)}=\frac{Pr(B\cap C_j\mid A)}{Pr(C_j\mid A)}=\frac{Pr(B\cap C_j\cap A)/Pr(A)}{Pr(C_j\cap A)/Pr(A)}=\frac{Pr(B\cap (C_j\cap A))}{Pr(C_j\cap A)}=Pr(B\mid C_j,A)$

So the law of total probability states that $Pr(B\mid A)=\sum _{j=1}^{n}Pr(B\mid C_j,A)Pr(C_j\mid A)$.

Now we see how the law of total probability interacts with conditional expectation. Let $X$ be a random variable. Then

$\mathbb{E}[X\mid A]={\sum }_{x}x\cdot Pr(X=x\mid A)={\sum }_{x}x{\sum }_{j}Pr(X=x\mid C_j,A)Pr(C_j\mid A)$

Here the event $X=x$ is playing the role of $B$ in the statement of the conditional law of total probability.

This is the basic trick of the proof; we keep conditioning on different things (actions, next states, rewards) and using the law of total probability.

By definition, $v_{\pi }(s)={\mathbb{E}}_{\pi }[G_t\mid S_t=s]$. Now rewrite $G_t=R_{t+1}+\gamma G_{t+1}$ and use the linearity of expectation to get ${\mathbb{E}}_{\pi }[R_{t+1}\mid S_t=s]+\gamma {\mathbb{E}}_{\pi }[G_{t+1}\mid S_t=s]$. From here, we can work separately with ${\mathbb{E}}_{\pi }[R_{t+1}\mid S_t=s]$ and ${\mathbb{E}}_{\pi }[G_{t+1}\mid S_t=s]$ for a while.

Using the law of total probability while conditioning over actions, we have

${\mathbb{E}}_{\pi }[R_{t+1}\mid S_t=s]={\sum }_{r}r\cdot Pr(R_{t+1}=r\mid S_t=s)={\sum }_{r}r{\sum }_{a}Pr(R_{t+1}=r\mid A_t=a,S_t=s)Pr(A_t=a\mid S_t=s)$

Using the convention that $Pr(A_t=a\mid S_t=s)=\pi (a\mid s)$, this becomes

${\mathbb{E}}_{\pi }[R_{t+1}\mid S_t=s]={\sum }_{r}r{\sum }_{a}Pr(R_{t+1}=r\mid A_t=a,S_t=s)\pi (a\mid s)$

Now we can reorder the sums to get

${\sum }_a\pi (a\mid s){\sum }_{r}rPr(R_{t+1}=r\mid A_t=a,S_t=s)$

Again, using the law of total probability (in its conjunction form) while conditioning this time over states, we have

${\sum }_a\pi (a\mid s){\sum }_{r}r{\sum }_{s^{\prime }}Pr(R_{t+1}=r,S_{t+1}=s^{\prime }\mid A_t=a,S_t=s)$

Reordering the sums, this becomes

${\sum }_a\pi (a\mid s){\sum }_{s^{\prime }}{\sum }_{r}Pr(R_{t+1}=r,S_{t+1}=s^{\prime }\mid A_t=a,S_t=s)r$

Sutton and Barto abbreviate $Pr(R_{t+1}=r,S_{t+1}=s^{\prime }\mid A_t=a,S_t=s)$ as $p(s^{\prime },r\mid s,a)$ (strictly speaking, we should track the timestep parameter $t$, but we will omit this detail here). We can also combine the nested sums ${\sum }_{s^{\prime }}{\sum }_r$ into a single sum that iterates over pairs $(s^{\prime },r)$. So we obtain

${\sum }_a\pi (a\mid s){\sum }_{s^{\prime },r}p(s^{\prime },r\mid s,a)r$

This completes the part for ${\mathbb{E}}_{\pi }[R_{t+1}\mid S_t=s]$. In other words, we have shown that

${\mathbb{E}}_{\pi }[R_{t+1}\mid S_t=s]={\sum }_a\pi (a\mid s){\sum }_{s^{\prime },r}p(s^{\prime },r\mid s,a)r$