Refactor lab

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

questions.md

@@ -1,13 +1,13 @@
 # Questions
 
-## BabySnake
-
-## Checkpoint 1: Before training
+## Checkpoint 1
 
 1. How do you decide where to move in BabySnake? Explain how to choose moves
    in enough detail that someone else could follow your instructions.
 
-2. How many distinct states are there for BabySnake? If we assume that all four 
+## Checkpoint 2
+
+2. How many distinct states are there for BabySnake on a 4×4 grid? If we assume that all four
    arrow keys are valid actions in every state, how many rows would the full Q-table contain?
 
 3. The discount factor γ (gamma) can range from 0 to 1. What would be the effect of setting
@@ -16,36 +16,13 @@
 4. The learning rate α (alpha) can also range from 0 to 1. What would be the effect of setting
    α to 0? What about 1?
 
-5. Calculate the new Q-value for the situation described. Explain your answer.
+5. Calculate the new Q-value for ((2, 2, 3, 3), RIGHT). Explain your answer.
 
-6. Implement `choose_action` and `update_q` in `q_learning.py`, then run
+## Checkpoint 3
 
-   ```
-   python test_q_learning.py
-   ```
+6. At what episode did the agent start reliably finding food?
 
-   Get every test passing before moving on — errors are much easier to track
-   down here than during training.
-
-
----
-
-## Checkpoint 2: After training
-
-Train your Q-learning agent to consistently score 3 or more food items per
-episode, then watch it play:
-
-```
-python train_babysnake.py
-```
-
-**At what episode did the agent start reliably finding food?**
-
-
-**Print `q_table` after training. Can you read the policy?** For a state you
-pick, does the highest Q-value point toward the food?
-
-
-**How does the trained agent's behavior compare to the reasoning you wrote
-down in Checkpoint 1?**
+7. Add `print(Q)` to `train_babysnake.py` before the `watch` call and run it again. Can you
+   read the policy? For a given state, does the highest Q-value point toward the food?
 
+8. How does the trained agent's behavior compare to the reasoning you wrote down in question 1?

test_q_learning.py

@@ -1,62 +0,0 @@
-# test_q_learning.py
-# ------------
-# Defines tests for `q_learning`. Run this program with `python test_q_learning.py`.
-# You don't need to edit this file.
-#
-# Get every test here passing before you run train_babysnake.py — errors are
-# much easier to spot here than during training.
-
-from unittest import TestCase, main
-import random
-
-from q_learning import choose_action, update_q
-
-ACTIONS = ["UP", "DOWN", "LEFT", "RIGHT"]
-
-
-class TestChooseAction(TestCase):
-    def test_greedy_picks_highest_q_value(self):
-        q_table = {("s", "UP"): 1.0, ("s", "DOWN"): 5.0, ("s", "LEFT"): 2.0, ("s", "RIGHT"): 0.0}
-        self.assertEqual(choose_action(q_table, "s", ACTIONS, epsilon=0.0), "DOWN")
-
-    def test_unseen_state_defaults_to_zero_and_picks_first_action(self):
-        self.assertEqual(choose_action({}, "new_state", ACTIONS, epsilon=0.0), ACTIONS[0])
-
-    def test_fully_random_explores_more_than_one_action(self):
-        random.seed(0)
-        q_table = {("s", "UP"): 100.0}  # UP is clearly the best action
-        results = {choose_action(q_table, "s", ACTIONS, epsilon=1.0) for _ in range(50)}
-        self.assertGreater(len(results), 1)
-
-    def test_always_returns_a_valid_action(self):
-        for _ in range(20):
-            result = choose_action({}, "s", ACTIONS, epsilon=0.5)
-            self.assertIn(result, ACTIONS)
-
-
-class TestUpdateQ(TestCase):
-    def test_basic_bellman_update(self):
-        q_table = {("s", "UP"): 0.0}
-        update_q(q_table, "s", "UP", reward=1.0, next_state="t", actions=ACTIONS, alpha=0.5, gamma=0.9)
-        # old_q=0, best_next_q=0 (unseen) -> target=1.0, new_q = 0 + 0.5*(1.0-0) = 0.5
-        self.assertAlmostEqual(q_table[("s", "UP")], 0.5)
-
-    def test_uses_best_next_q_value(self):
-        q_table = {("s", "UP"): 0.0, ("t", "UP"): 2.0, ("t", "DOWN"): 5.0}
-        update_q(q_table, "s", "UP", reward=0.0, next_state="t", actions=ACTIONS, alpha=1.0, gamma=1.0)
-        # target = 0 + 1.0*5.0 = 5.0; alpha=1 fully replaces the old value
-        self.assertAlmostEqual(q_table[("s", "UP")], 5.0)
-
-    def test_alpha_zero_means_no_change(self):
-        q_table = {("s", "UP"): 3.0}
-        update_q(q_table, "s", "UP", reward=10.0, next_state="t", actions=ACTIONS, alpha=0.0, gamma=0.9)
-        self.assertAlmostEqual(q_table[("s", "UP")], 3.0)
-
-    def test_only_updates_the_given_state_action_pair(self):
-        q_table = {}
-        update_q(q_table, "s", "UP", reward=1.0, next_state="t", actions=ACTIONS, alpha=0.5, gamma=0.9)
-        self.assertEqual(set(q_table.keys()), {("s", "UP")})
-
-
-if __name__ == '__main__':
-    main()