65 lines
2.2 KiB
Python
65 lines
2.2 KiB
Python
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()
|