Why isn't the implementation of my policy evaluation for a simple MDP converging?

Question

I am trying to code out a policy evaluation algorithm to find the $V^\pi(s)$ for all states. The following diagram below shows the MDP.

In this case i let p = q = 0.5. the rewards for each states are independent of action. I.e $r(\sigma_0)$ = $r(\sigma_2)$ = 0,$r(\sigma_1)$ = 1, $r(\sigma_3)$ = 10. Terminal state is $r(\sigma_3)$

I have the following policy, {0:1, 1:0, 2:0}, where key is the state and value is the action. 0 for $a_0$ and 1 for $a_1$.

#Policy Iteration solver for FUN
class PolicyEvaluation:
    def __init__(self, policies):
        self.N = 3
        self.pi = policies
        self.actions = [0, 1] # a0 and a1
        self.discount = 0.7
        self.states = [i for i in range(self.N + 1)]


    def terminalState(self, state):
        return state == 3

    # assume p = q = 0.5
    def succProbReward(self, state):
        # (newState, probability, reward)
        spr_list = []
        if (state == 0 and self.pi[state] == 0):
            spr_list.append([1, 1.0, 1])
        elif (state == 0 and self.pi[state] == 1):
            spr_list.append([2, 1.0, 0])
        elif (state == 1 and self.pi[state] == 0):
            spr_list.append([2, 0.5, 0])
            spr_list.append([0, 0.5, 0])
        elif (state == 2 and self.pi[state] == 0):
            spr_list.append([1, 1.0, 0])
        elif (state == 2 and self.pi[state] == 1):
            spr_list.append([3, 0.5, 10])
            spr_list.append([2, 0.5, 0])
        return spr_list


def policyEvaluation(mdp):
    # initialize
    V = {} 
    for state in mdp.states:
        V[state] = 0

    def V_pi(state):
        return sum(prob * (reward + mdp.discount*V[newState]) for prob, reward, newState in
        mdp.succProbReward(state))

    while True:
    # compute new values (newV) given old values (V)
        newV = {}
        for state in mdp.states:
            if mdp.terminalState(state):
                newV[state] = 0
            else:
                newV[state] = V_pi(state)

        if max(abs(V[state] - newV[state]) for state in mdp.states) < 1e-10:
            break
        V = newV
        print(V)
    print(V)



pE = PolicyEvaluation({0:1, 1:0, 2:0})
print(pE.states)
print(pE.succProbReward(0))
policyIteration(pE)

I've tried to run the code above to find the values for each state, however, I am not converging with my values.

Is there something wrong that I did?

Answer 1

The issue is that in your list comprehension in def V_pi(state) you have

return sum(prob * (reward + mdp.discount*V[newState]) for prob, reward, newState in
        mdp.succProbReward(state))

whereas with the way you have defined the succProbReward output, it should be

return sum(prob * (reward + mdp.discount*V[newState]) for newState, prob, reward in
        mdp.succProbReward(state))

When I run this it converges immediately with a reward of 0 for all states, which I believe is correct for the policy you specified. If I change the policy it also seems to give reasonable results.