Unit 14: Reinforcement Learning

# ==============================================================================
# DevAscent Academy - Unit 14: Reinforcement Learning
# Standalone Tabular Q-Learning Engine on Cliff Walking GridWorld
# Pure Python Implementation - Zero External Dependencies
# ==============================================================================

import random

class CliffGridWorld:
    """A 4x12 GridWorld environment with a hazardous cliff along the bottom."""
    def __init__(self, width=12, height=4):
        self.width = width
        self.height = height
        self.start_state = (0, 3) # Bottom-left corner (x, y)
        self.goal_state = (11, 3) # Bottom-right corner (x, y)
        # The cliff spans coordinates (1, 3) through (10, 3)
        self.cliff = [(x, 3) for x in range(1, 11)]
        self.agent_pos = self.start_state

    def reset(self):
        self.agent_pos = self.start_state
        return self.agent_pos

    def step(self, action):
        """
        Executes action: 0=UP, 1=RIGHT, 2=DOWN, 3=LEFT.
        Returns: (next_state, reward, is_terminal)
        """
        x, y = self.agent_pos
        if action == 0:   y = max(0, y - 1) # Move Up
        elif action == 1: x = min(self.width - 1, x + 1) # Move Right
        elif action == 2: y = min(self.height - 1, y + 1) # Move Down
        elif action == 3: x = max(0, x - 1) # Move Left

        next_state = (x, y)
        self.agent_pos = next_state

        # Check Cliff Hazard
        if next_state in self.cliff:
            # Fatal drop into cliff: -100 penalty, respawn at start
            return self.start_state, -100, False

        # Check Goal Destination
        if next_state == self.goal_state:
            return next_state, +20, True # Reached goal successfully

        # Standard step cost to encourage path efficiency
        return next_state, -1, False


class QLearningAgent:
    """Tabular Q-Learning agent with Epsilon-Greedy exploration."""
    def __init__(self, actions=4, lr=0.1, gamma=0.95, epsilon=1.0, epsilon_decay=0.995, min_epsilon=0.02):
        self.actions = actions
        self.lr = lr
        self.gamma = gamma
        self.epsilon = epsilon
        self.epsilon_decay = epsilon_decay
        self.min_epsilon = min_epsilon
        # Q-Table dictionary: { (x, y): [Q_up, Q_right, Q_down, Q_left] }
        self.q_table = {}

    def get_q_values(self, state):
        if state not in self.q_table:
            # Initialize unseen states with zero value
            self.q_table[state] = [0.0 for _ in range(self.actions)]
        return self.q_table[state]

    def choose_action(self, state):
        """Epsilon-Greedy Action Selection Policy."""
        q_vals = self.get_q_values(state)
        # Explore: pick random action
        if random.random() < self.epsilon:
            return random.randint(0, self.actions - 1)
        # Exploit: pick action with maximum Q-value
        max_val = max(q_vals)
        best_actions = [i for i, v in enumerate(q_vals) if v == max_val]
        return random.choice(best_actions)

    def learn(self, state, action, reward, next_state, done):
        """Temporal Difference (TD) Q-Value Update."""
        old_q = self.get_q_values(state)[action]

        if done:
            target = reward
        else:
            next_max_q = max(self.get_q_values(next_state))
            target = reward + (self.gamma * next_max_q)

        # Apply Bellman learning delta
        td_error = target - old_q
        self.q_table[state][action] = old_q + (self.lr * td_error)

    def decay_exploration(self):
        """Smoothly reduces epsilon after each episode."""
        if self.epsilon > self.min_epsilon:
            self.epsilon *= self.epsilon_decay


# ==============================================================================
# TRAINING & CONVERGENCE SIMULATION
# ==============================================================================
if __name__ == "__main__":
    env = CliffGridWorld()
    agent = QLearningAgent(lr=0.1, gamma=0.95, epsilon=1.0, epsilon_decay=0.992)

    num_episodes = 500
    print(f"--- Commencing Q-Learning Training across {num_episodes} Episodes ---")

    for episode in range(1, num_episodes + 1):
        state = env.reset()
        total_reward = 0
        steps = 0

        while True:
            steps += 1
            action = agent.choose_action(state)
            next_state, reward, done = env.step(action)
            agent.learn(state, action, reward, next_state, done)

            state = next_state
            total_reward += reward

            if done or steps > 150: # Safety step cap
                break

        agent.decay_exploration()

        if episode % 100 == 0 or episode == 1:
            print(f"Episode {episode:03d} | Total Reward: {total_reward:4d} | Steps: {steps:3d} | Epsilon: {agent.epsilon:.3f}")

    # --- Demonstrate Final Optimal Greedy Policy ---
    print("\n================ EVALUATING OPTIMAL LEARNED POLICY ================")
    agent.epsilon = 0.0 # Strict exploitation
    state = env.reset()
    path = [state]
    action_names = ["UP", "RIGHT", "DOWN", "LEFT"]
    steps = 0

    while state != env.goal_state and steps < 20:
        steps += 1
        act = agent.choose_action(state)
        next_s, rew, done = env.step(act)
        path.append(next_s)
        state = next_s

    print(f"Optimal Trajectory Discovered ({len(path)} waypoints):")
    print(" -> ".join([str(p) for p in path]))
    print("Safety Verification: Zero cliff drops encountered on converged path.")