# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 2
#     language: python
#     name: python2
# ---

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

env = GridworldEnv()

env.P.keys()
#for saTuple in env.P[0]:
#    print env.P[0][saTuple]
print env.P[0]
print np.ones([env.nS, env.nA]) / env.nA


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 (prob, next_state, reward, done) tuple.
        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:
    # prevent infinite loops
    for cnt in range(1000):
        delta = 0.0
        for si in range(policy.shape[0]):
            v = V[si]
            tmp = 0
            for saInd in env.P[si]:
                saTuple = env.P[si][saInd][0]
                tmp += policy[si][saInd] * saTuple[0] * (saTuple[2] + discount_factor * V[saTuple[1]])
            V[si] = tmp
            delta = max(delta,abs(v - V[si]))
        if delta < theta:
            break
    return np.array(V)


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

# 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)



StarcoderdataJupyterScriptsDedupFiltered