通过代码学Sutton强化学习第四章动态规划

MDP = Dict[State, Dict[Action, List[Tuple[Prob, State, Reward, bool]]]]

# P[state][action] = [
#    (prob1, next_state1, reward1, is_done),
#    (prob2, next_state2, reward2, is_done), ...]

class Action(Enum):
    UP = 0
    DOWN = 1
    LEFT = 2
    RIGHT = 3

State = int
Reward = float
Prob = float
Policy = Dict[State, Dict[Action, Prob]]
Value = List[float]
StateSet = Set[int]
NonTerminalStateSet = Set[int]
MDP = Dict[State, Dict[Action, List[Tuple[Prob, State, Reward, bool]]]]
# P[s][a] = [(prob, next_state, reward, is_done), ...]

class GridWorldEnv(discrete.DiscreteEnv):
    """
    Grid World environment described in Sutton and Barto Reinforcement Learning 2nd, chapter 4.
    """

    def __init__(self, shape=[4,4]):
        self.shape = shape
        nS = np.prod(shape)
        nA = len(list(Action))
        MAX_R = shape[0]
        MAX_C = shape[1]
        self.grid = np.arange(nS).reshape(shape)
        isd = np.ones(nS) / nS

        # P[s][a] = [(prob, next_state, reward, is_done), ...]
        P: MDP = {}
        action_delta = {Action.UP: (-1, 0), Action.DOWN: (1, 0), Action.LEFT: (0, -1), Action.RIGHT: (0, 1)}
        for s in range(0, MAX_R * MAX_C):
            P[s] = {a.value : [] for a in list(Action)}
            is_terminal = self.is_terminal(s)
            if is_terminal:
                for a in list(Action):
                    P[s][a.value] = [(1.0, s, 0, True)]
            else:
                r = s // MAX_R
                c = s % MAX_R
                for a in list(Action):
                    neighbor_r = min(MAX_R-1, max(0, r + action_delta[a][0]))
                    neighbor_c = min(MAX_C-1, max(0, c + action_delta[a][1]))
                    s_ = neighbor_r * MAX_R + neighbor_c
                    P[s][a.value] = [(1.0, s_, -1, False)]

        super(GridWorldEnv, self).__init__(nS, nA, P, isd)

def policy_evaluate(policy: Policy, env: GridWorldEnv, gamma=1.0, theta=0.0001):
    V = np.zeros(env.nS)
    while True:
        delta = 0
        for s in range(env.nS):
            v = 0
            for a, action_prob in enumerate(policy[s]):
                for prob, next_state, reward, done in env.P[s][a]:
                    v += action_prob * prob * (reward + gamma * V[next_state])
            delta = max(delta, np.abs(v - V[s]))
            V[s] = v
        if delta < theta:
            break
    return np.array(V)

[[  0.         -13.99931242 -19.99901152 -21.99891199]
 [-13.99931242 -17.99915625 -19.99908389 -19.99909436]
 [-19.99901152 -19.99908389 -17.99922697 -13.99942284]
 [-21.99891199 -19.99909436 -13.99942284   0.        ]]

def action_value(env: GridWorldEnv, state: State, V: StateValue, gamma=1.0) -> ActionValue:
    q = np.zeros(env.nA)
    for a in range(env.nA):
        for prob, next_state, reward, done in env.P[state][a]:
            q[a] += prob * (reward + gamma * V[next_state])
    return q

def policy_improvement(env: GridWorldEnv, policy: Policy, V: StateValue, gamma=1.0) -> bool:
    policy_stable = True

    for s in range(env.nS):
        old_action = np.argmax(policy[s])
        Q_s = action_value(env, s, V)
        best_action = np.argmax(Q_s)
        policy[s] = np.eye(env.nA)[best_action]

        if old_action != best_action:
            policy_stable = False
    return policy_stable


def policy_iteration(env: GridWorldEnv, policy: Policy, gamma=1.0) -> Tuple[Policy, StateValue]:
    iter = 0
    while True:
        V = policy_evaluate(policy, env, gamma)
        policy_stable = policy_improvement(env, policy, V)
        iter += 1

        if policy_stable:
            return policy, V

def value_iteration(env:GridWorldEnv, gamma=1.0, theta=0.0001) -> Tuple[Policy, StateValue]:
    V = np.zeros(env.nS)
    while True:
        delta = 0
        for s in range(env.nS):
            action_values = action_value(env, s, V, gamma=gamma)
            best_action_value = np.max(action_values)
            delta = max(delta, np.abs(best_action_value - V[s]))
            V[s] = best_action_value
        if delta < theta:
            break

    policy = np.zeros([env.nS, env.nA])
    for s in range(env.nS):
        action_values = action_value(env, s, V, gamma=gamma)
        best_action = np.argmax(action_values)
        policy[s, best_action] = 1.0

    return policy, V

def build_reverse_mapping(env:GridWorldEnv) -> Dict[State, Set[State]]:
    MAX_R, MAX_C = env.shape[0], env.shape[1]
    mapping = {s: set() for s in range(0, MAX_R * MAX_C)}
    action_delta = {Action.UP: (-1, 0), Action.DOWN: (1, 0), Action.LEFT: (0, -1), Action.RIGHT: (0, 1)}
    for s in range(0, MAX_R * MAX_C):
        r = s // MAX_R
        c = s % MAX_R
        for a in list(Action):
            neighbor_r = min(MAX_R - 1, max(0, r + action_delta[a][0]))
            neighbor_c = min(MAX_C - 1, max(0, c + action_delta[a][1]))
            s_ = neighbor_r * MAX_R + neighbor_c
            mapping[s_].add(s)
    return mapping

def value_iteration_async(env:GridWorldEnv, gamma=1.0, theta=0.0001) -> Tuple[Policy, StateValue]:
    mapping = build_reverse_mapping(env)

    V = np.zeros(env.nS)
    changed_state_set = set(s for s in range(env.nS))

    iter = 0
    while len(changed_state_set) > 0:
        changed_state_set_ = set()
        for s in changed_state_set:
            action_values = action_value(env, s, V, gamma=gamma)
            best_action_value = np.max(action_values)
            v_diff = np.abs(best_action_value - V[s])
            if v_diff > theta:
                changed_state_set_.update(mapping[s])
                V[s] = best_action_value
        changed_state_set = changed_state_set_
        iter += 1

    policy = np.zeros([env.nS, env.nA])
    for s in range(env.nS):
        action_values = action_value(env, s, V, gamma=gamma)
        best_action = np.argmax(action_values)
        policy[s, best_action] = 1.0

    return policy, V