[Python] R_max algorithm doesn't converge to the right policy

I have a task where I need to implement an R_max algorithm with modified policy itieration over the forzen lake problem.

I first tried to just implement a regular modified policy iteration and it worked but when I tried to adjust it to r max requirments (using fake model with a fake state) it doesn't converge to the right policy.

I have tried to debug but I am not even sure where to start and chatGPT wasnt a great help.

Unfortuntely I am also not sure where excatly the problem is. I would really appreciate if you could have a look over my code and maybe spot a few mistakes.

sorry for the long code but I really dont know where the problem is to try to focus it to a more relevant section

import numpy as np

class Rmax:
def __init__(self, env, r_max=2.5, gamma=0.95, k=20, num_iterations=5, epsilon=1e-6, iterations=30000):
# init model basic data structure
self.env = env
self.num_states = env.observation_space.n + 1
self.num_actions = env.action_space.n
self.r_max = r_max
self.gamma = gamma
self.k = k
self.num_iterations = num_iterations
self.epsilon = epsilon
self.iterations = iterations
#init fake Tr table
self.Tr = np.ones((self.num_states, self.num_actions, self.num_states))
for state in range(self.num_states-1):
for action in range(self.num_actions):
for next_state in range(self.num_states - 1):
self.Tr[state][action][next_state] = 1/(self.num_states-1) # we assume transition to all real states in uniform distribution

self.rewards = np.zeros((self.num_states, self.num_actions))
for action in range(self.num_actions):
self.rewards[self.num_states-1][action] = self.r_max # reward for fake state is 1 all other rewards are 0

#counters
self.state_action_counter = np.zeros((self.num_states, self.num_actions), dtype=int)
self.state_action_state_counter = np.zeros((self.num_states, self.num_actions, self.num_states), dtype=int)

#known table
self.known = np.zeros((self.num_states, self.num_actions), dtype=bool)

self.policy = np.zeros(self.num_states, dtype=int)

#and 0 for all values
self.values = np.ones(self.num_states)


def train(self):
# Training loop is Solve model discover new states and solve again
for i in range(self.iterations):
print(f'iteration: {i}/{self.iterations}')
self.solve_model()


def solve_model(self):
# solving the model with MPI requires to steps policy evaluation and policy improvement
self.policy_eval()
stable = self.policy_improvement()
# now we run policy till we know
changed = False
done, truncated = False, False
while not changed and not done and not truncated:
state = self.env.reset()[0]
action = self.policy[state]
next_state, reward, done, truncated, info = self.env.step(action)
self.state_action_counter[state][action] += 1
self.state_action_state_counter[state][action][next_state] += 1
if self.state_action_counter[state][action] >= self.k:
self.known[state][action] = True
changed = True
self.update_model()

def update_model(self):
for state in range(self.num_states):
for action in range(self.num_actions):
if self.known[state][action]:
for next_state in range(self.num_states):
if self.state_action_state_counter[state][action][next_state] != 0:
self.Tr[state][action][next_state] = (
self.state_action_state_counter[state][action][next_state] / self.state_action_counter[state][action])
else:
self.Tr[state][action] = np.ones(self.num_states) / self.num_states



def policy_eval(self):
# in MPI we do a limited number of policy evaluation
for _ in range(self.num_iterations):
new_values = self.values.copy()
for state in range(self.num_states):
action = self.policy[state]
new_value = 0
for next_state in range(self.num_states):
prob = self.Tr[state][action][next_state]
reward = self.rewards[state][action]
new_value += prob *(reward + self.gamma * self.values[next_state])
new_values[state] = new_value


#also if policy converged we can stop
if np.max(np.abs(new_values - self.values)) < self.epsilon:
self.values = new_values
break
self.values = new_values

def policy_improvement(self):
stable = True
for state in range(self.num_states):
old_action = self.policy[state]
action_values = np.zeros( self.num_actions)
for action in range(self.num_actions):
for next_state in range(self.num_states):
prob = self.Tr[state][action][next_state]
reward = self.rewards[state][action]
action_values[action] += prob * (reward + self.gamma * self.values[next_state])
best_action = np.argmax(action_values)
self.policy[state] = best_action
if old_action != best_action:
stable = False
return stable


I tried to play with the values go over the algorithm again and getting help from chatGPT unfortunalty didnt able to identify the problem

Continue reading...