Tic-Tac-Toe Reinforcement Learning: Q-Learning Agent Training and Gameplay

Someone

12 min read

·

Nov 30, 2024

--

Listen

Share

In this blog post, we will walk through the implementation of a Q-learning agent that learns to play Tic-Tac-Toe. Using reinforcement learning (RL), specifically the Q-learning algorithm, we’ll create an agent capable of optimizing its moves over time.

import time
import pygame
import random
import numpy as np
from collections import defaultdict

class TicTacToe:
    def __init__(self) -> None:
        """
            Args: 
                None

            Returns:
                None

            Concept:
                Initializes class Tictactoe and sets up the board, current player, and player wins and pygame.
        """
        
        self.board = np.zeros((3, 3), dtype=int)
        self.current_player = 1
        self.player1_wincount = 0
        self.player2_wincount = 0
        self.draw_count = 0
        self.init_pygame()

    def init_pygame(self) -> None:
        """
            Args:
                None

            Returns:
                None

            Concept:
                Initializes pygame and sets up the screen, font, and clock.
        """
        
        pygame.init()
        self.screen = pygame.display.set_mode((300, 300))
        pygame.display.set_caption("Tic-Tac-Toe Training")
        self.font = pygame.font.Font(None, 74)
        self.clock = pygame.time.Clock()

    def reset(self) -> np.ndarray:
        """
            Args:
                None
            
            Returns:
                None
            
            Concept:
                Resets the board, current player, and player wins.
        """
        
        self.board = np.zeros((3, 3), dtype=int)
        self.current_player = 1
        self.render()
        
        return self.board

    def available_actions(self) -> list:
        """
            Args:
                None
            
            Returns:
                list of available actions
            
            Concept:
                Returns a list of all available actions for the agent on the board.
        """
        
        return list(zip(*np.where(self.board == 0)))
    
    def make_move(self, action: tuple[int, int]) -> tuple[int, bool]:
        """
            Args:
                action: Action selected by the agent (tuple)
            
            Returns:
                value of reward and whether the game is over with this move
                
            Concept:
                Makes a move on the board and updates the current player.
        """
        
        if self.board[action] == 0:
            self.board[action] = self.current_player
            reward, done = self.check_winner()
            self.render()
            self.current_player = -self.current_player
            
            return reward, done
        return 0, False
    
    def check_winner(self) -> tuple[int, bool]:
        """
            Args:
                None
            
            Returns:
                player who won and that the game is over
            
            Concept:
                Checks if there is a winner on the board.
        """
        
        for i in range(3):
            if np.all(self.board[i, :] == self.current_player) or np.all(self.board[:, i] == self.current_player):
                print(f"Player {1 if self.current_player == 1 else 2} wins!")
                print(self.board)
                
                if self.current_player == 1:
                    self.player1_wincount += 1
                else:
                    self.player2_wincount += 1
                
                return 1 if self.current_player == 1 else -1, True
        
        if np.all(np.diag(self.board) == self.current_player) or np.all(np.diag(np.fliplr(self.board)) == self.current_player):
            print(f"Player {1 if self.current_player == 1 else 2} wins!")
            print(self.board)
            
            if self.current_player == 1:
                self.player1_wincount += 1
            else:
                self.player2_wincount += 1
                
            return 1 if self.current_player == 1 else -1, True
        
        if not np.any(self.board == 0):
            print("It's a draw!")
            print(self.board)
            
            self.draw_count += 1
            
            return 0, True
        
        return 0, False

    def check_critical(self, player: int) -> tuple[float, tuple[int, int]]:
        """

        Args:
            player: specifier of player

        Returns:
            reward and critical position on the board and location of critical position.

        Concept:
            Checks if there is a critical position on the board. Critical position is a position that if one move can decide the winner.
        """
        
        for i in range(3):
            if np.sum(self.board[i, :] == player) == 2 and np.any(self.board[i, :] == 0):
                return 0.1, (i, np.where(self.board[i, :] == 0)[0][0])
            
            if np.sum(self.board[:, i] == player) == 2 and np.any(self.board[:, i] == 0):
                return 0.1, (np.where(self.board[:, i] == 0)[0][0], i)
        
        if np.sum(np.diag(self.board) == player) == 2 and np.any(np.diag(self.board) == 0):
            idx = np.where(np.diag(self.board) == 0)[0][0]
            return 0.1, (idx, idx)
        
        if np.sum(np.diag(np.fliplr(self.board)) == player) == 2 and np.any(np.diag(np.fliplr(self.board)) == 0):
            idx = np.where(np.diag(np.fliplr(self.board)) == 0)[0][0]
            return 0.1, (idx, 2 - idx)
        
        opponent = -player
        for i in range(3):
            if np.sum(self.board[i, :] == opponent) == 2 and np.any(self.board[i, :] == 0):
                return -0.1, (i, np.where(self.board[i, :] == 0)[0][0])
            
            if np.sum(self.board[:, i] == opponent) == 2 and np.any(self.board[:, i] == 0):
                return -0.1, (np.where(self.board[:, i] == 0)[0][0], i)
        
        if np.sum(np.diag(self.board) == opponent) == 2 and np.any(np.diag(self.board) == 0):
            idx = np.where(np.diag(self.board) == 0)[0][0]
            return -0.1, (idx, idx)
        
        if np.sum(np.diag(np.fliplr(self.board)) == opponent) == 2 and np.any(np.diag(np.fliplr(self.board)) == 0):
            idx = np.where(np.diag(np.fliplr(self.board)) == 0)[0][0]
            return -0.1, (idx, 2 - idx)
        
        return 0.2, None

    def render(self) -> None:
        """
            Args:
                None
            
            Returns:
                None
            
            Concept:
                Renders the board on the screen.
        """
        
        self.screen.fill((255, 255, 255))
        for row in range(3):
            for col in range(3):
                if self.board[row, col] == 1:
                    pygame.draw.line(self.screen, (0, 0, 0), (col * 100 + 15, row * 100 + 15), (col * 100 + 85, row * 100 + 85), 15)
                    pygame.draw.line(self.screen, (0, 0, 0), (col * 100 + 15, row * 100 + 85), (col * 100 + 85, row * 100 + 15), 15)
                elif self.board[row, col] == -1:
                    pygame.draw.circle(self.screen, (0, 0, 0), (col * 100 + 50, row * 100 + 50), 40, 15)
        
        for i in range(1, 3):
            pygame.draw.line(self.screen, (0, 0, 0), (0, i * 100), (300, i * 100), 5)
            pygame.draw.line(self.screen, (0, 0, 0), (i * 100, 0), (i * 100, 300), 5)

        pygame.display.flip()
        
        # time.sleep(0.1)  # Delay for 100 milliseconds # Uncomment to see the moves taken by the agent in real time

    def close_pygame(self) -> None:
        """
            Args:
                None
            
            Returns:
                None
            
            Concept:
                Closes the pygame window.
        """
        
        pygame.quit()

class QLearningAgent:
    def __init__(self, player: int, epsilon: int=0.1, alpha: int=0.5, gamma: int=0.9) -> None:
        """
            Args:
                player: The specifier for the current player
                epsilon: The exploration rate
                alpha: The learning rate
                gamma: The discount factor
            
            Returns:
                None
            
            Concept:
                Initializes the class QLearningAgent with the given parameters; and sets the q_table to an empty dictionary.
        """
        
        self.q_table = defaultdict(lambda: 0)
        self.epsilon = epsilon
        self.alpha = alpha
        self.gamma = gamma
        self.player = player
    
    def get_state(self, board: np.ndarray) -> tuple:
        """
            Args:
                board: The current board being played at the moment
            
            Returns:
                State as the board condition
                
            Concept:
                Returns the state as the board condition.
        """
        
        return tuple(map(tuple, board))
    
    def choose_action(self, board: np.ndarray, available_actions: list) -> int:
        """
            Args:
                board: The current board being played at the moment
                available_actions: A list of actions that can be taken on the current board
            
            Returns:
                A calculated action of the available actions
            
            Concept:
                Returns a calculated action of the available actions and if epsilon value is greater than a random value then a random action is returned.
        """
        
        if random.uniform(0, 1) < self.epsilon:
            return random.choice(available_actions)
        
        q_values = [self.q_table[(self.get_state(board), action)] for action in available_actions]
        max_q = max(q_values)
        
        return random.choice([a for a, q in zip(available_actions, q_values) if q == max_q])
    
    def update_q_table(self, state: tuple, action: int, reward: int, next_state: tuple, next_actions: list) -> None:
        """
            Args:
                state: The current board at the moment
                action: Action taken by the agent
                reward: reward given for that action in that condition
                next_state: The next board to played
                next_actions: Actions available to model on the next board
            
            Returns:
                None
            
            Concept:
                Updates the q_table with the given parameters.
        """
        
        current_q = self.q_table[(state, action)]
        if next_actions:
            max_next_q = max([self.q_table[(next_state, a)] for a in next_actions])
        else:
            max_next_q = 0
        self.q_table[(state, action)] = current_q + self.alpha * (reward + self.gamma * max_next_q - current_q)
    
    def save_model(self, filename: str) -> None:
        """
            Args:
                filename: filename for the q_table to be saved in
            
            Returns:
                None
            
            Concept:
                Saves the q_table in the given filename.
        """
        
        import pickle
        with open(filename, 'wb') as f:
            pickle.dump(dict(self.q_table), f)
    
    def load_model(self, filename: str) -> None:
        """
            Args:
                filename: filename for the q_table to be loaded from
            
            Returns:
                None
            
            Concept:
                Loads the q_table from the given filename.
        """
        
        import pickle
        with open(filename, 'rb') as f:
            self.q_table = defaultdict(lambda: 0, pickle.load(f))

def train(episodes: int, epsilon: float, alpha: float, gamma: float) -> None:
    env = TicTacToe()
    agent1 = QLearningAgent(player=1, epsilon=epsilon, alpha=alpha, gamma=gamma)
    agent2 = QLearningAgent(player=-1, epsilon=epsilon, alpha=alpha, gamma=gamma)
    
    try:
        for episode in range(episodes):
            print(f"Episode {episode + 1}/{episodes}")
            
            state = env.reset()
            done = False
            while not done:
                player = env.current_player
                agent = agent1 if player == 1 else agent2
                state_tuple = agent.get_state(state)
                actions = env.available_actions()
                action = agent.choose_action(state, actions)
                
                reward, critical_location = env.check_critical(player)
                
                """
                    If critical location is given and agent wins, give 2 reward points
                    If critical location is given and agent doesn't mark there and loses, give 2 penalty points
                """
                if critical_location:
                    if action == critical_location:
                        reward = 2
                    elif done and reward == -1:
                        reward = -2
                
                reward_move, done = env.make_move(action)
                next_state_tuple = agent.get_state(env.board)
                next_actions = env.available_actions()
                agent.update_q_table(state_tuple, action, reward + reward_move, next_state_tuple, next_actions)
                
                if done and reward != 0:
                    other_agent = agent2 if player == 1 else agent1
                    other_agent.update_q_table(state_tuple, action, -reward, next_state_tuple, next_actions)
                
                for event in pygame.event.get():
                    if event.type == pygame.QUIT:
                        pygame.quit()
                        return

    finally:
        env.close_pygame()
        print(f"Player 1 wins: {env.player1_wincount} times")
        print(f"Player 2 wins: {env.player2_wincount} times")
        print(f"Draws: {env.draw_count}")

    agent1.save_model('Files/agent1_q_table.pkl')
    agent2.save_model('Files/agent2_q_table.pkl')

if __name__ == "__main__":
    epsilon = 0.25    # Exploration Rate
    alpha = 0.05      # Learning Rate
    gamma = 0.15       # Discount Factor
    episodes = 10000  # Iteration Count
    
    train(episodes=episodes, epsilon=epsilon, alpha=alpha, gamma=gamma)

import time
import pygame
import pickle
import random

import numpy as np
from collections import defaultdict

class TicTacToe:
    def __init__(self) -> None:
        self.board = np.zeros((3, 3), dtype=int)
        self.current_player = 1
        self.init_pygame()

    def init_pygame(self) -> None:
        pygame.init()
        self.screen = pygame.display.set_mode((300, 300))
        pygame.display.set_caption("Tic-Tac-Toe Testing")
        self.font = pygame.font.Font(None, 74)
        self.clock = pygame.time.Clock()

    def reset(self) -> np.ndarray:
        self.board = np.zeros((3, 3), dtype=int)
        self.current_player = 1
        self.render()
        
        return self.board

    def available_actions(self) -> list:
        return list(zip(*np.where(self.board == 0)))

    def make_move(self, action) -> bool:
        if self.board[action] == 0:
            self.board[action] = self.current_player
            done, winner = self.check_winner()
            self.render()
            self.current_player = -self.current_player
            
            return done, "Player 1 wins!" if winner == 1 else "Player 2 wins!" if winner == -1 else "Draw!" if winner == 0 else None
        return False, None

    def check_winner(self) -> bool:
        for i in range(3):
            if np.all(self.board[i, :] == self.current_player) or np.all(self.board[:, i] == self.current_player):
                print(f"Player {1 if self.current_player == 1 else 2} wins!")
                return True, 1 if self.current_player == 1 else -1
        
        if np.all(np.diag(self.board) == self.current_player) or np.all(np.diag(np.fliplr(self.board)) == self.current_player):
            print(f"Player {1 if self.current_player == 1 else 2} wins!")
            return True, 1 if self.current_player == 1 else -1
        
        if not np.any(self.board == 0):
            print("It's a draw!")
            return True, 0
        
        return False, None

    def render(self) -> None:
        self.screen.fill((255, 255, 255))
        for row in range(3):
            for col in range(3):
                if self.board[row, col] == 1:
                    pygame.draw.line(self.screen, (0, 0, 0), (col * 100 + 15, row * 100 + 15), (col * 100 + 85, row * 100 + 85), 15)
                    pygame.draw.line(self.screen, (0, 0, 0), (col * 100 + 15, row * 100 + 85), (col * 100 + 85, row * 100 + 15), 15)
                elif self.board[row, col] == -1:
                    pygame.draw.circle(self.screen, (0, 0, 0), (col * 100 + 50, row * 100 + 50), 40, 15)
        
        for i in range(1, 3):
            pygame.draw.line(self.screen, (0, 0, 0), (0, i * 100), (300, i * 100), 5)
            pygame.draw.line(self.screen, (0, 0, 0), (i * 100, 0), (i * 100, 300), 5)

        pygame.display.flip()

    def close_pygame(self) -> None:
        pygame.quit()

class QLearningAgent:
    def __init__(self, player: int, q_table: dict) -> None:
        self.q_table = q_table
        self.player = player
        self.actions_taken = ""
    
    def get_state(self, board: np.ndarray) -> tuple:
        return tuple(map(tuple, board))
    
    def choose_action(self, board: np.ndarray, available_actions: list) -> tuple:
        q_values = [self.q_table[(self.get_state(board), action)] for action in available_actions]
        max_q = max(q_values)
        chosen_action = random.choice([a for a, q in zip(available_actions, q_values) if q == max_q])
        
        self.actions_taken += f"""\nGame State: \n{board}\n"""
        self.actions_taken += f"""Available Actions: \n"""
        for q_value in q_values:
            self.actions_taken += f"\t{q_value}\n"
        self.actions_taken += f"""Chosen Action: \n\t{max_q}\n"""
        
        print(f"\nPlayer {1 if self.player == 1 else 2} Q-values: {q_values}")
        print(f"\nPlayer {1 if self.player == 1 else 2} chooses action: {chosen_action} \nwith Q-value: {max_q}\n")
        
        return chosen_action
    
    def save_actions_taken(self, filename: str) -> None:
        with open(filename, 'w') as f:
            f.write(str(self.actions_taken))

def load_q_table(filename: str) -> dict:
    with open(filename, 'rb') as f:
        q_table = pickle.load(f)
    
    return defaultdict(lambda: 0, q_table)

def user_move(available_actions: list) -> tuple:
    while True:
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                pygame.quit()
                exit()
            if event.type == pygame.MOUSEBUTTONDOWN:
                x, y = event.pos
                row, col = y // 100, x // 100
                if (row, col) in available_actions:
                    return (row, col)

def main() -> None:
    agent1_file = "Files/agent1_q_table.pkl"
    agent2_file = "Files/agent2_q_table.pkl"
    actions_taken_file = "Files/actions_taken.txt"
    
    env = TicTacToe()
    
    user_player = int(input("Choose player 1 (X) or player 2 (O): "))
    if user_player == 1:
        user_player = 1
    else:
        user_player = -1
    
    agent_player = -user_player
    agent_q_table_file = agent1_file if agent_player == 1 else agent2_file
    
    q_table = load_q_table(agent_q_table_file)
    agent = QLearningAgent(player=agent_player, q_table=q_table)
    
    while True:
        print("Starting new game...")
        state = env.reset()
        done = False
        agent.actions_taken += "Game Start: \n"
        
        while not done:
            if env.current_player == user_player:
                available_actions = env.available_actions()
                action = user_move(available_actions)
            else:
                available_actions = env.available_actions()
                action = agent.choose_action(state, available_actions)
            
            print(env.board)
            
            done, game_result = env.make_move(action)

            if done:
                print()
                time.sleep(2)
        
        agent.actions_taken += f"Game End: {game_result}\n\n\n"
        agent.save_actions_taken(actions_taken_file)

if __name__ == "__main__":
    main()

We’ll cover the core components, from defining the game environment to implementing the learning agent, and finally, testing the trained agent in gameplay.

Overview of Q-Learning and Tic-Tac-Toe

Q-learning is a model-free reinforcement learning algorithm that enables an agent to learn optimal actions by interacting with an environment. The agent updates a Q-table, where each entry represents the expected future reward for a given state-action pair.

In the case of Tic-Tac-Toe, the agent will learn to make moves that maximize its chances of winning, by exploring various game states, updating its strategy, and exploiting its knowledge to make better decisions.

1. Implementing the Game Environment: The TicTacToe Class

First, let’s define the game environment. The TicTacToe class is where we implement the logic of the game, including board states, available actions, and rewards.

Importing Required Libraries

import time
import pygame
import random

import numpy as np
from collections import defaultdict

We use the pygame library for graphical display, numpy for handling the game board as an array, and random for randomizing moves. The defaultdict from collections helps us store the Q-table, initializing unknown states with a default value of 0.

Defining the Game Class

class TicTacToe:
    def __init__(self) -> None:
        self.board = np.zeros((3, 3), dtype=int)  # 3x3 grid for Tic-Tac-Toe
        self.current_player = 1  # Player 1 starts
        self.winner = None  # No winner at the start

    def reset(self):
        """Resets the game state."""
        self.board = np.zeros((3, 3), dtype=int)
        self.current_player = 1
        self.winner = None
    
    def available_actions(self):
        """Returns available positions on the board."""
        return list(zip(*np.where(self.board == 0)))
    
    def play(self, action):
        """Makes a move on the board and updates the game state."""
        x, y = action
        if self.board[x, y] == 0:
            self.board[x, y] = self.current_player
            self.check_winner()
            self.current_player = 3 - self.current_player  # Switch player
            return True
        return False
    
    def check_winner(self):
        """Checks for a winner in the game."""
        for row in self.board:
            if np.all(row == 1) or np.all(row == 2):
                self.winner = self.current_player
                return
        for col in self.board.T:
            if np.all(col == 1) or np.all(col == 2):
                self.winner = self.current_player
                return
        if np.all(np.diag(self.board) == 1) or np.all(np.diag(self.board) == 2):
            self.winner = self.current_player
            return
        if np.all(np.diag(np.fliplr(self.board)) == 1) or np.all(np.diag(np.fliplr(self.board)) == 2):
            self.winner = self.current_player
            return

This class encapsulates the core game logic, including resetting the board, making a move, checking for available actions, and determining if there’s a winner. The reset method starts a new game, while available_actions identifies valid positions for the next move.

2. Implementing the Q-learning Agent: The QLearningAgent Class

The heart of this project is the Q-learning agent. The QLearningAgent class implements the Q-learning algorithm and interacts with the TicTacToe environment to improve its gameplay over time.

Defining the QLearningAgent Class

class QLearningAgent:
    def __init__(self, player: int, epsilon: float = 0.1, alpha: float = 0.5, gamma: float = 0.9) -> None:
        self.q_table = defaultdict(lambda: 0)  # Initialize Q-table
        self.epsilon = epsilon  # Exploration factor
        self.alpha = alpha  # Learning rate
        self.gamma = gamma  # Discount factor
        self.player = player  # Which player the agent is (1 or 2)
        self.actions_taken = []  # Store actions for analysis

The QLearningAgent class initializes the Q-table, which stores the expected rewards for each state-action pair. We also define the exploration-exploitation trade-off with the epsilon parameter, the learning rate (alpha), and the discount factor (gamma).

Choosing an Action: Exploration vs Exploitation

def choose_action(self, state, available_actions):
    """Choose an action based on exploration vs. exploitation."""
    if random.uniform(0, 1) < self.epsilon:
        # Exploration: Choose a random action
        return random.choice(available_actions)
    else:
        # Exploitation: Choose the action with the highest Q-value
        q_values = [self.q_table[(tuple(state.flatten()), action)] for action in available_actions]
        max_q_value = max(q_values)
        best_actions = [action for action, q in zip(available_actions, q_values) if q == max_q_value]
        return random.choice(best_actions)

In the choose_action method, the agent either explores random moves or exploits its learned knowledge by selecting the action with the highest Q-value.

Updating the Q-Table: Q-learning Update Rule

def update_q_table(self, state, action, reward, next_state, available_actions):
    """Update the Q-table using the Q-learning update rule."""
    future_q_values = [self.q_table[(tuple(next_state.flatten()), next_action)] for next_action in available_actions]
    max_future_q = max(future_q_values) if future_q_values else 0
    current_q = self.q_table[(tuple(state.flatten()), action)]
    new_q = current_q + self.alpha * (reward + self.gamma * max_future_q - current_q)
    self.q_table[(tuple(state.flatten()), action)] = new_q

The update_q_table method updates the Q-value of the chosen action based on the reward received and the maximum future reward the agent expects from the next state.

3. Training the Agent: train Function

The train function handles the iterative process of training the agent by running multiple episodes of Tic-Tac-Toe.

Training Loop

def train(episodes: int, epsilon: float, alpha: float, gamma: float) -> None:
    agent = QLearningAgent(player=1, epsilon=epsilon, alpha=alpha, gamma=gamma)
    env = TicTacToe()

    for episode in range(episodes):
        print(f"Episode {episode + 1}/{episodes}")
        state = env.reset()
        done = False
        while not done:
            available_actions = env.available_actions()
            action = agent.choose_action(state, available_actions)
            env.play(action)
            next_state = env.board
            reward = 1 if env.winner == 1 else (-1 if env.winner == 2 else 0)
            agent.update_q_table(state, action, reward, next_state, available_actions)
            if env.winner:
                done = True
            state = next_state

In this loop, the agent interacts with the environment and learns over several episodes, updating its Q-table after each move. The training stops when a winner is determined.

4. Testing the Trained Agent: main Function

Once the agent has been trained, we can test it in a real game scenario against either a user or another agent.

Main Game Loop

def main() -> None:
    agent = QLearningAgent(player=2)
    env = TicTacToe()
    # Load trained Q-table (if available)
    while True:
        print("Starting new game…")
        state = env.reset()
        while not done:
            if env.current_player == 1:  # Human player's turn
                available_actions = env.available_actions()
                action = user_move(available_actions)
                env.play(action)
            else:  # Agent's turn
                available_actions = env.available_actions()
                action = agent.choose_action(state, available_actions)
                env.play(action)
                print(f"Agent chooses action: {action}")
            
            if env.winner:
                print(f"Player {env.winner} wins!")
                break
            state = env.board

The main function handles user interaction. The user makes moves by calling the user_move function, while the agent makes its moves based on the learned Q-values.

5. User Interaction: user_move Function

To allow the user to play against the trained agent, we implement the user_move function that handles user input.

def user_move(available_actions: list) -> tuple:
    """Handle user input for making a move."""
    # Code for capturing user input and returning the selected move

Conclusion

This blog post demonstrates a basic implementation of a Q-learning agent trained to play Tic-Tac-Toe. The agent learns effective strategies by interacting with the game environment, improving its decision-making over time. While this implementation is simple, it showcases the

power of reinforcement learning in game-playing agents. Future improvements might include more advanced reward functions, optimizations for larger state spaces, or testing against multiple agents to explore different strategies.

Feel free to explore the code on GitHub and follow me on Twitter for more updates and projects!