from nis import match
from turtle import st
import numpy as np
high = 1
low = 0
search, wait, recharge = 0, 1, 2
def print_policy(pi):
def action(i):
return{
0: 'search',
1: 'wait',
2: 'recharge'
}.get(i)
print('low ->', action(pi[0]))
print('high ->', action(pi[1]))
def mdp(alpha, beta, rs, rw):
S = (low, high) #0: low, 1:high
A = (search, wait, recharge)
p = np.zeros((len(S), len(S), len(A)))
p[high, high, search] = alpha
p[high, low, search] = 1 - alpha
p[high, high, wait] = 1
p[low, low, search] = beta
p[low, high, search] = 1 - beta
p[low, low, wait] = 1
p[low, high, recharge] = 1
r = np.zeros((len(S), len(S), len(A)))
r[high, high, search] = rs
r[high, low, search] = rs
r[high, high, wait] = rw
r[low, low, search] = rs
r[low, high, search] = -3
r[low, low, wait] = rw
r[low, high, recharge] = 0
return S, A, p, r
def policy_iteration(mdp, gamma, epsilon):
S, A, p, r = mdp
#1. Initialization
v = np.zeros(len(S))
pi = np.zeros(len(S))
while(True):
#2. Policy Evaluation
while(True):
delta = 0
for s in S:
v_current = v[s]
#Update v[s]:
v[s] = sum([p[int(s), int(sp), int(pi[int(s)])]
* (r[int(s), int(sp), int(pi[int(s)])] + gamma * v[sp]) for sp in S])
delta = max(delta, abs(v_current - v[s]))
if delta < epsilon * (1-gamma)/(2*gamma):
break
#3. Policy Improvement
stable = True
for s in S:
q_best = v[s]
for a in A:
q = sum([p[int(s), int(sp), int(a)]
* (r[int(s), int(sp), int(a)] + gamma * v[sp]) for sp in S])
if q > q_best:
pi[s] = a
q_best = q
stable = False
if stable:
return pi
alpha = .5
beta = .5
rs = 10.
rw = 2.
mdp = mdp(alpha, beta, rs, rw)
pi_star = policy_iteration(mdp, gamma=.9, epsilon=.0001)
print_policy(pi_star)