In [46]:
import numpy as np
import sys
if "../" not in sys.path:
  sys.path.append("../") 
from lib.envs.gridworld import GridworldEnv

In [47]:
env = GridworldEnv()

In [49]:
def policy_eval(policy, env, discount_factor=1.0, theta=0.00001):
    """
    Evaluate a policy given an environment and a full description of the environment's dynamics.
    
    Args:
        policy: [S, A] shaped matrix representing the policy.
        env: OpenAI env. env.P represents the transition probabilities of the environment.
            env.P[s][a] is a list of transition tuples (prob, next_state, reward, done).
        theta: We stop evaluation once our value function change is less than theta for all states.
        discount_factor: gamma discount factor.
    
    Returns:
        Vector of length env.nS representing the value function.
    """
    # Start with a random (all 0) value function
    V = np.zeros(env.nS)
    while True:
        delta = 0
        for s in range(env.nS):
            # expected future reward when taking action a from state s
            state_value = 0
            for a in range(env.nA):
                state_action_value = 0
                for prob, next_state, reward, done in env.P[s][a]:
                    # bellman expectation update
                    state_action_value += prob * (reward + discount_factor * V[next_state])
                state_value += policy[s][a] * state_action_value
            
            # calculate delta, then update state value
            delta = max(delta, abs(V[s] - state_value))
            V[s] = state_value
        
        # delta lower than threshold
        if delta < theta:
            break
    return np.array(V)

In [50]:
random_policy = np.ones([env.nS, env.nA]) / env.nA
v = policy_eval(random_policy, env)

In [51]:
# Test: Make sure the evaluated policy is what we expected
expected_v = np.array([0, -14, -20, -22, -14, -18, -20, -20, -20, -20, -18, -14, -22, -20, -14, 0])
np.testing.assert_array_almost_equal(v, expected_v, decimal=2)

In [ ]: