Refactor lab

q_learning/__init__.py

@@ -0,0 +1,99 @@
+import random
+from itertools import count
+from retro_gamer import GameEnvironment, GameMetadata
+
+class QLearning:
+    """
+    Trains a policy to play a game using Q-learning. 
+    The Q-table maintains an estimate of the quality of every action at every state, 
+    and the estimates are improved through training. 
+
+    Arguments:
+        episodes (int): How many episodes to run.
+        alpha (float): Learning rate.
+        gamma (float): Discount factor.
+        epsilon (float): Starting exploration rate.
+        epsilon_decay (float): Multiply epsilon by this each episode.
+        epsilon_min (float): Epsilon never falls below this.
+        max_steps_per_episode (int): Safety cutoff. Without this, a lucky
+            random walk that keeps finding food (each pickup restores more
+            energy than a turn costs) can make an episode run far longer
+            than intended, or even effectively forever.
+    """
+
+    def __init__(
+        self, episodes=1000, alpha=0.1, gamma=0.95, epsilon=1.0, 
+        epsilon_decay=0.995, epsilon_min=0.05, max_steps_per_episode=500
+    ):
+        self.episodes = episodes
+        self.alpha = alpha
+        self.gamma = gamma
+        self.epsilon = epsilon
+        self.epsilon_decay = epsilon_decay
+        self.epsilon_min = epsilon_min
+        self.max_steps_per_episode = max_steps_per_episode
+
+    def train(self, env, actions):
+        """Trains the policy by updating estimates of the quality of state/actions
+        in the Q-table.
+
+        Arguments:
+            env (retro_gamer.GameEnvironment): a game environment, which allows
+                stepping through a game one action at a time.
+        """
+        self.Q = {}
+        for episode in range(self.episodes):
+            state = env.reset()
+            turn = 0
+            total_reward = 0
+            for turn in range(self.max_steps_per_episode):
+                action = self.choose_action(state, actions)
+                next_state, reward, done = env.step(action)
+                self.update_q(state, action, reward, next_state, actions)
+                state = next_state
+                total_reward += reward
+            self.epsilon = max(self.epsilon_min, self.epsilon * self.epsilon_decay)
+            self.report_progress(episode, total_reward)
+        return self.Q
+
+    def choose_action(self, state, actions):
+        """Choose an action using an epsilon-greedy policy.
+        With probability `self.epsilon`, return a random action from `actions`.
+        Otherwise, return the action with the highest Q-value in `self.Q`
+        for the given `state`. If a (state, action) pair has not been seen
+        before, treat its Q-value as 0.
+        """
+        if random.random() < self.epsilon:
+            return random.choice(actions)
+        else:
+            action_qualities = [(self.Q.get((state, a), 0), a) for a in actions]
+            best_q, best_action = sorted(action_qualities, reverse=True)[0]
+            return best_action
+
+    def update_q(self, state, action, reward, next_state, next_actions):
+        """Update an entry in self.Q. The update rule is: 
+
+            Q(s, a) <- Q(s, a) + alpha * (r + gamma * max_a' Q(s', a') - Q(s, a))
+
+        where:
+            s, a                the state we were in and the action we took
+            r                   the reward we received
+            s'                  the next state 
+            max_a' Q(s', a')    the best possible Q-value of actions from the next state
+            alpha               the learning rate (use self.alpha)
+            gamma               the discount factor (use self.gamma)
+        """
+        future_qs = [(self.Q.get((next_state, a), 0), a) for a in next_actions]
+        best_future_q, best_future_action = sorted(future_qs, reverse=True)[0]
+        q = self.Q.get((state, action), 0)
+        self.Q[(state, action)] = q + self.alpha * (reward + self.gamma * best_future_q - q)
+
+    def report_progress(self, episode, total_reward):
+        if (episode + 1) % 100 == 0:
+            print(
+                f"Episode {episode + 1:5d}  "
+                f"reward={total_reward:6.1f}  "
+                f"epsilon={self.epsilon:.3f}  "
+                f"q_entries={len(self.Q)}"
+            )
+

q_learning/tests.py

@@ -0,0 +1,64 @@
+from unittest import TestCase, main
+import random
+
+from q_learning import QLearning
+
+ACTIONS = ["UP", "DOWN", "LEFT", "RIGHT"]
+
+
+class TestChooseAction(TestCase):
+    def test_greedy_picks_highest_q_value(self):
+        q = QLearning(epsilon=0.0)
+        q.Q = {("s", "UP"): 1.0, ("s", "DOWN"): 5.0, ("s", "LEFT"): 2.0, ("s", "RIGHT"): 0.0}
+        self.assertEqual(q.choose_action("s", ACTIONS), "DOWN")
+
+    def test_unseen_state_defaults_to_zero_and_picks_first_action(self):
+        q = QLearning(epsilon=0.0)
+        q.Q = {}
+        self.assertEqual(q.choose_action("new_state", ACTIONS), ACTIONS[0])
+
+    def test_fully_random_explores_more_than_one_action(self):
+        random.seed(0)
+        q = QLearning(epsilon=1.0)
+        q.Q = {("s", "UP"): 100.0}
+        results = {q.choose_action("s", ACTIONS) for _ in range(50)}
+        self.assertGreater(len(results), 1)
+
+    def test_always_returns_a_valid_action(self):
+        q = QLearning(epsilon=0.5)
+        q.Q = {}
+        for _ in range(20):
+            result = q.choose_action("s", ACTIONS)
+            self.assertIn(result, ACTIONS)
+
+
+class TestUpdateQ(TestCase):
+    def test_basic_bellman_update(self):
+        q = QLearning(alpha=0.5, gamma=0.9)
+        q.Q = {("s", "UP"): 0.0}
+        q.update_q("s", "UP", 1.0, "t", ACTIONS)
+        # old_q=0, best_next_q=0 (unseen) -> target=1.0, new_q = 0 + 0.5*(1.0-0) = 0.5
+        self.assertAlmostEqual(q.Q[("s", "UP")], 0.5)
+
+    def test_uses_best_next_q_value(self):
+        q = QLearning(alpha=1.0, gamma=1.0)
+        q.Q = {("s", "UP"): 0.0, ("t", "UP"): 2.0, ("t", "DOWN"): 5.0}
+        q.update_q("s", "UP", 0.0, "t", ACTIONS)
+        # target = 0 + 1.0*5.0 = 5.0; alpha=1 fully replaces the old value
+        self.assertAlmostEqual(q.Q[("s", "UP")], 5.0)
+
+    def test_alpha_zero_means_no_change(self):
+        q = QLearning(alpha=0.0, gamma=0.9)
+        q.Q = {("s", "UP"): 3.0}
+        q.update_q("s", "UP", 10.0, "t", ACTIONS)
+        self.assertAlmostEqual(q.Q[("s", "UP")], 3.0)
+
+    def test_only_updates_the_given_state_action_pair(self):
+        q = QLearning(alpha=0.5, gamma=0.9)
+        q.Q = {}
+        q.update_q("s", "UP", 1.0, "t", ACTIONS)
+        self.assertEqual(set(q.Q.keys()), {("s", "UP")})
+
+
+if __name__ == '__main__':
+    main()