function value_iteration(states, actions, transition, reward, gamma=0.9):
   value = {s: 0 for s in states}

   repeat until value stops changing:
      for s in states:
        value[s] = max(reward(s, a) + gamma * value[transition(s, a)] for a in
 actions)
   policy = {s: argmax(actions, key=lambda a: reward(s, a) + gamma * value[transit
 ion(s, a)]) for s in states}
   return policy
