Add Q-Learning algorithm implementation with epsilon-greedy policy an… (#13402)

* Add Q-Learning algorithm implementation with epsilon-greedy policy and grid world demo

* Add Q-Learning algorithm implementation with epsilon-greedy policy and grid world demo code

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* bug fixes and linting

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* cls

* bug fix and hints

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: idklol22

machine_learning/q_learning.py

@@ -0,0 +1,206 @@
+"""
+Q-Learning is a widely-used model-free algorithm in reinforcement learning that
+learns the optimal action-value function Q(s, a), which tells an agent the expected
+utility of taking action a in state s and then following the optimal policy after.
+It is able to find the best policy for any given finite Markov decision process (MDP)
+without requiring a model of the environment.
+
+See: [https://en.wikipedia.org/wiki/Q-learning](https://en.wikipedia.org/wiki/Q-learning)
+"""
+
+import random
+from collections import defaultdict
+
+# Type alias for state
+type State = tuple[int, int]
+
+# Hyperparameters for Q-Learning
+LEARNING_RATE = 0.1
+DISCOUNT_FACTOR = 0.97
+EPSILON = 0.2
+EPSILON_DECAY = 0.995
+EPSILON_MIN = 0.01
+
+# Global Q-table to store state-action values
+q_table: dict[State, dict[int, float]] = defaultdict(lambda: defaultdict(float))
+
+# Environment variables for simple grid world
+SIZE = 4
+GOAL = (SIZE - 1, SIZE - 1)
+current_state = (0, 0)
+
+
+def get_q_value(state: State, action: int) -> float:
+    """
+    Get Q-value for a given state-action pair.
+
+    >>> q_table.clear()
+    >>> get_q_value((0, 0), 2)
+    0.0
+    """
+    return q_table[state][action]
+
+
+def get_best_action(state: State, available_actions: list[int]) -> int:
+    """
+    Get the action with maximum Q-value in the given state.
+
+    >>> q_table.clear()
+    >>> q_table[(0, 0)][1] = 0.7
+    >>> q_table[(0, 0)][2] = 0.7
+    >>> q_table[(0, 0)][3] = 0.5
+    >>> get_best_action((0, 0), [1, 2, 3]) in [1, 2]
+    True
+    """
+    if not available_actions:
+        raise ValueError("No available actions provided")
+    max_q = max(q_table[state][a] for a in available_actions)
+    best = [a for a in available_actions if q_table[state][a] == max_q]
+    return random.choice(best)
+
+
+def choose_action(state: State, available_actions: list[int]) -> int:
+    """
+    Choose action using epsilon-greedy policy.
+
+    >>> q_table.clear()
+    >>> old_epsilon = EPSILON
+    >>> EPSILON = 0.0
+    >>> q_table[(0, 0)][1] = 1.0
+    >>> q_table[(0, 0)][2] = 0.5
+    >>> result = choose_action((0, 0), [1, 2])
+    >>> EPSILON = old_epsilon  # Restore
+    >>> result
+    1
+    """
+    global EPSILON
+    if not available_actions:
+        raise ValueError("No available actions provided")
+    if random.random() < EPSILON:
+        return random.choice(available_actions)
+    return get_best_action(state, available_actions)
+
+
+def update(
+    state: State,
+    action: int,
+    reward: float,
+    next_state: State,
+    next_available_actions: list[int],
+    done: bool = False,
+    alpha: float | None = None,
+    gamma: float | None = None,
+) -> None:
+    """
+    Perform Q-value update for a transition using the Q-learning rule.
+
+    Q(s,a) <- Q(s,a) + alpha * (r + gamma * max_a' Q(s',a') - Q(s,a))
+
+    >>> q_table.clear()
+    >>> update((0, 0), 1, 1.0, (0, 1), [1, 2], done=True, alpha=0.5, gamma=0.9)
+    >>> get_q_value((0, 0), 1)
+    0.5
+    """
+    global LEARNING_RATE, DISCOUNT_FACTOR
+    alpha = alpha if alpha is not None else LEARNING_RATE
+    gamma = gamma if gamma is not None else DISCOUNT_FACTOR
+    max_q_next = (
+        0.0
+        if done or not next_available_actions
+        else max(get_q_value(next_state, a) for a in next_available_actions)
+    )
+    old_q = get_q_value(state, action)
+    new_q = (1 - alpha) * old_q + alpha * (reward + gamma * max_q_next)
+    q_table[state][action] = new_q
+
+
+def get_policy() -> dict[State, int]:
+    """
+    Extract a deterministic policy from the Q-table.
+
+
+    >>> q_table.clear()
+    >>> q_table[(1, 2)][1] = 2.0
+    >>> q_table[(1, 2)][2] = 1.0
+    >>> get_policy()[(1, 2)]
+    1
+    """
+    policy: dict[State, int] = {}
+    for s, a_dict in q_table.items():
+        if a_dict:
+            policy[s] = max(a_dict, key=lambda a: a_dict[a])
+    return policy
+
+
+def reset_env() -> State:
+    """
+    Reset the environment to initial state.
+
+    >>> old_state = current_state
+    >>> current_state = (1, 1)  # Simulate non-initial state
+    >>> result = reset_env()
+    >>> current_state = old_state  # Restore for other tests
+    >>> result
+    (0, 0)
+    """
+    global current_state
+    current_state = (0, 0)
+    return current_state
+
+
+def get_available_actions_env() -> list[int]:
+    """
+    Get available actions in the current environment state.
+    """
+    return [0, 1, 2, 3]  # 0: up, 1: right, 2: down, 3: left
+
+
+def step_env(action: int) -> tuple[State, float, bool]:
+    """
+    Take a step in the environment with the given action.
+    """
+    global current_state
+    x, y = current_state
+    if action == 0:  # up
+        x = max(0, x - 1)
+    elif action == 1:  # right
+        y = min(SIZE - 1, y + 1)
+    elif action == 2:  # down
+        x = min(SIZE - 1, x + 1)
+    elif action == 3:  # left
+        y = max(0, y - 1)
+    next_state = (x, y)
+    reward = 10.0 if next_state == GOAL else -1.0
+    done = next_state == GOAL
+    current_state = next_state
+    return next_state, reward, done
+
+
+def run_q_learning() -> None:
+    """
+    Run Q-Learning on the simple grid world environment.
+    """
+    global EPSILON
+    episodes = 200
+    for _ in range(episodes):
+        state = reset_env()
+        done = False
+        while not done:
+            actions = get_available_actions_env()
+            action = choose_action(state, actions)
+            next_state, reward, done = step_env(action)
+            next_actions = get_available_actions_env()
+            update(state, action, reward, next_state, next_actions, done)
+            state = next_state
+        EPSILON = max(EPSILON * EPSILON_DECAY, EPSILON_MIN)
+    policy = get_policy()
+    print("Learned Policy (state: action):")
+    for s, a in sorted(policy.items()):
+        print(f"{s}: {a}")
+
+
+if __name__ == "__main__":
+    import doctest
+
+    doctest.testmod()
+    run_q_learning()