def policy_evaluation(policy, V, mdp, k=20):
    """Calculate the value function following a policy"""
    R, T, gamma = mdp.R, mdp.T, mdp.gamma
    # For horizon k
    for i in range(k):
        # For each state
        for s in mdp.states:
            V[s] = R(s) + gamma * sum([p * V[s] for (p, s1) in T(s, policy[s])])
    return V