# What is the optimal policy for this MDP?

What is the optimal policy for this MDP?

Since this is model-based RL you can use policy iteration method of dynamic programming, and since this MDP is very simple you can intuitively guess an optimal policy for your initial iterated policy $$\pi_0$$ as $$\pi_0(a_1|s_1)=1, \pi_0(a_2|s_2)=1$$, where $$a_1$$ is the switch action and $$a_2$$ is the stay action.

In the first policy evaluation step we need to solve for the Bellman equation of both states:

$$v_{\pi_0}(s_1)=\sum_a\pi_0(a|s_1)\big[\sum_rp(r|s_1,a_1)r+\gamma \sum_{s'}p(s'|s_1,a_1)v_{\pi_0}(s')\big]=0+0.5(0.2v_{\pi_0}(s_1)+0.8v_{\pi_0}(s_2))$$

$$v_{\pi_0}(s_2)=\sum_a\pi_0(a|s_2)\big[\sum_rp(r|s_2,a_2)r+\gamma \sum_{s'}p(s'|s_2,a_2)v_{\pi_0}(s')\big]=1+0.5v_{\pi_0}(s_2)$$

This is a simple system of linear equations of 2 variables $$v_{\pi_0}(s_1),v_{\pi_0}(s_2)$$, you can easily check the solutions are $$v_{\pi_0}(s_1)=\frac{8}{9},v_{\pi_0}(s_2)=2$$.

Then in the first policy improvement step we need to compute action value for each state-action pair:

$$q_{\pi_0}(s_1,a_1)=\sum_rp(r|s_1,a_1)r+\gamma \sum_{s'}p(s'|s_1,a_1)v_{\pi_0}(s')=0+0.5(0.2v_{\pi_0}(s_1)+0.8v_{\pi_0}(s_2))=\frac{8}{9}$$

$$q_{\pi_0}(s_1,a_2)=\sum_rp(r|s_1,a_2)r+\gamma \sum_{s'}p(s'|s_1,a_2)v_{\pi_0}(s')=0+0.5v_{\pi_0}(s_1)=\frac{4}{9}$$

$$q_{\pi_0}(s_2,a_1)=\sum_rp(r|s_2,a_1)r+\gamma \sum_{s'}p(s'|s_2,a_1)v_{\pi_0}(s')=0+0.5(0.2v_{\pi_0}(s_2)+0.8v_{\pi_0}(s_1))=\frac{5}{9}$$

$$q_{\pi_0}(s_2,a_2)=\sum_rp(r|s_2,a_2)r+\gamma \sum_{s'}p(s'|s_2,a_2)v_{\pi_0}(s')=1+0.5v_{\pi_0}(s_2)=2$$

Therefore the improved policy is updated as $$\pi_1(a_1|s_1)=1,\pi_1(a_2|s_2)=1$$, which apparently already converges to $$\pi_0$$ initialized in the above first section. So we can conclude an optimal policy for this MDP is $$\pi^*(a_1|s_1)=1, \pi^*(a_2|s_2)=1$$.