human
2. (1.7)
Implement a MDP corresponding to the one in question 3 of HW4 from [AI-edX]. Modify the value iteration algorithm to consider rewards that depend on $(s, a, s')$, where $s$, $a$, and $s'$ correspond to the current state, the action and the next state respectively."
Please write code following the instructions in jupyter notebook style.
assistant
#Here is code to 2.

class ExtendedMDP(MDP):

    def __init__(self, transition_matrix, rewards, terminals, init, gamma=.9):
        # All possible actions.
        actlist = []
        for state in transition_matrix.keys():
            actlist.extend(transition_matrix[state])
        actlist = list(set(actlist))

        MDP.__init__(self, init, actlist, terminals=terminals, gamma=gamma)
        self.t = transition_matrix
        self.reward = rewards
        for state in self.t:
            self.states.add(state)

    def T(self, state, action):
        return [(prob, new_state) for new_state, prob 
                in self.t[state][action].items()]

    def R(self, state, action, statep):
        "Returns a numeric reward for the combination the current state, the action and the next state."
        return self.reward[state][action][statep]
    
def ext_value_iteration(mdp, epsilon=0.001):
    "Solving an MDP by value iteration. Uses a reward function that depends on s, a and s1 "
    U1 = {s: 0 for s in mdp.states}
    R, T, gamma = mdp.R, mdp.T, mdp.gamma
    while True:
        U = U1.copy()
        delta = 0
        for s in mdp.states:
            # Your code here #
            U1[s] = 
            ##################
            delta = max(delta, abs(U1[s] - U[s]))
        if delta < epsilon * (1 - gamma) / gamma:
            return U
        
def ext_calculate_v_star(gamma):
    '''
    This function must create an instance of ExtendedMDP that corresponds to the 
    MDP in question 3 of HW4 and use it to calculate the expected reward for each state.
    The function receives as parameter the value of gamma. 
    The function must return a dictionary with the V* values. 
    '''
    # Your code here #
    t = 
    rewards = 
    ##################
    emdp = ExtendedMDP(t, rewards, [], None, gamma=gamma)
    v_star = ext_value_iteration(emdp, epsilon = 0.00001)
    return v_star