Update README.md

README.md

@@ -33,7 +33,6 @@ import numpy as np
 import pickle
 
 def greedy_policy(Qtable, state):
-    # Exploitation: take the action with the highest state, action value
     action = np.argmax(Qtable[state, :])
     return action
 
@@ -50,23 +49,24 @@ def evaluate_agent(env: gym.Env, max_steps: int, n_eval_episodes: int, Q: np.nda
     for episode in tqdm(range(n_eval_episodes)):
         if seed:
             state, info = env.reset(seed=seed[episode])
-    else:
-        state, info = env.reset()
-    step = 0
-    truncated = False
-    terminated = False
-    total_rewards_ep = 0
+        else:
+            state, info = env.reset()
 
-    for step in range(max_steps):
-        # Take the action (index) that have the maximum expected future reward given that state
-        action = greedy_policy(Q, state)
-        new_state, reward, terminated, truncated, info = env.step(action)
-        total_rewards_ep += reward
+        truncated = False
+        terminated = False
+        total_rewards_ep = 0
+
+        for step in range(max_steps):
+            action = greedy_policy(Q, state)
+            new_state, reward, terminated, truncated, info = env.step(action)
+            total_rewards_ep += reward
+
+            if terminated or truncated:
+                break
+            state = new_state
+
+        episode_rewards.append(total_rewards_ep)
 
-        if terminated or truncated:
-            break
-        state = new_state
-    episode_rewards.append(total_rewards_ep)
     mean_reward = np.mean(episode_rewards)
     std_reward = np.std(episode_rewards)
 
@@ -74,7 +74,7 @@ def evaluate_agent(env: gym.Env, max_steps: int, n_eval_episodes: int, Q: np.nda
 
 if __name__ == "__main__":
     file_path = load_from_hub(repo_id="BobChuang/q-Taxi-v1-5x5", filename="q-learning.pkl")
-    with open(file_path, 'rb') as f:
+    with open(file_path, "rb") as f:
         model = pickle.load(f)
 
     env = gym.make(model["env_id"], render_mode="rgb_array")