connect-four-MCTS.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Connect 4"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from collections import Counter\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "NUM_COLUMNS = 7\n",
    "COLUMN_HEIGHT = 6\n",
    "FOUR = 4\n",
    "\n",
    "# Board can be initialized with `board = np.zeros((NUM_COLUMNS, COLUMN_HEIGHT), dtype=np.byte)`\n",
    "# Notez Bien: Connect 4 \"columns\" are actually NumPy \"rows\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Basic Functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "def valid_moves(board):\n",
    "    \"\"\"Returns columns where a disc may be played\"\"\"\n",
    "    return [n for n in range(NUM_COLUMNS) if board[n, COLUMN_HEIGHT - 1] == 0]\n",
    "\n",
    "\n",
    "def play(board, column, player):\n",
    "    \"\"\"Updates `board` as `player` drops a disc in `column`\"\"\"\n",
    "    (index,) = next((i for i, v in np.ndenumerate(board[column]) if v == 0))\n",
    "    board[column, index] = player\n",
    "\n",
    "\n",
    "def take_back(board, column):\n",
    "    \"\"\"Updates `board` removing top disc from `column`\"\"\"\n",
    "    (index,) = [i for i, v in np.ndenumerate(board[column]) if v != 0][-1]\n",
    "    board[column, index] = 0\n",
    "\n",
    "\n",
    "def four_in_a_row(board, player):\n",
    "    \"\"\"Checks if `player` has a 4-piece line\"\"\"\n",
    "    return (\n",
    "        any(\n",
    "            all(board[c, r] == player)\n",
    "            for c in range(NUM_COLUMNS)\n",
    "            for r in (list(range(n, n + FOUR)) for n in range(COLUMN_HEIGHT - FOUR + 1))\n",
    "        )\n",
    "        or any(\n",
    "            all(board[c, r] == player)\n",
    "            for r in range(COLUMN_HEIGHT)\n",
    "            for c in (list(range(n, n + FOUR)) for n in range(NUM_COLUMNS - FOUR + 1))\n",
    "        )\n",
    "        or any(\n",
    "            np.all(board[diag] == player)\n",
    "            for diag in (\n",
    "                (range(ro, ro + FOUR), range(co, co + FOUR))\n",
    "                for ro in range(0, NUM_COLUMNS - FOUR + 1)\n",
    "                for co in range(0, COLUMN_HEIGHT - FOUR + 1)\n",
    "            )\n",
    "        )\n",
    "        or any(\n",
    "            np.all(board[diag] == player)\n",
    "            for diag in (\n",
    "                (range(ro, ro + FOUR), range(co + FOUR - 1, co - 1, -1))\n",
    "                for ro in range(0, NUM_COLUMNS - FOUR + 1)\n",
    "                for co in range(0, COLUMN_HEIGHT - FOUR + 1)\n",
    "            )\n",
    "        )\n",
    "    )\n",
    "\n",
    "def _mc(board, player):\n",
    "    p = -player\n",
    "    while valid_moves(board):\n",
    "        p = -p\n",
    "        c = np.random.choice(valid_moves(board))\n",
    "        play(board, c, p)\n",
    "        if four_in_a_row(board, p):\n",
    "            return p\n",
    "    return 0\n",
    "\n",
    "\n",
    "def montecarlo(board, player):\n",
    "    montecarlo_samples = 100\n",
    "    cnt = Counter(_mc(np.copy(board), player) for _ in range(montecarlo_samples))\n",
    "    return (cnt[1] - cnt[-1]) / montecarlo_samples\n",
    "\n",
    "\n",
    "def eval_board(board, player):\n",
    "    if four_in_a_row(board, 1):\n",
    "        # Alice won\n",
    "        return 1\n",
    "    elif four_in_a_row(board, -1):\n",
    "        # Bob won\n",
    "        return -1\n",
    "    else:\n",
    "        # Not terminal, let's simulate...\n",
    "        return montecarlo(board, player)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Author: Fra98\n",
    "### Four-Connect solver with Montecarlo tree exploration\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "PLAYERS = {1: \"A\", -1: \"B\"}\n",
    "MAX_ROUNDS = NUM_COLUMNS * COLUMN_HEIGHT\n",
    "NUM_ITERATIONS = 1000"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Auxiliary Functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "def initialize_board():\n",
    "    return np.zeros((NUM_COLUMNS, COLUMN_HEIGHT), dtype=np.byte)\n",
    "    \n",
    "def display(board):\n",
    "    for j in reversed(range(COLUMN_HEIGHT)):\n",
    "        for i in range(NUM_COLUMNS):\n",
    "            cell = board[i][j]\n",
    "            if cell == 1:\n",
    "                print(PLAYERS[1], end=\" \")\n",
    "            elif cell == -1:\n",
    "                print(PLAYERS[-1], end=\" \")\n",
    "            else:\n",
    "                print(\"-\", end=\" \")\n",
    "        print()\n",
    "\n",
    "\n",
    "def round_number(board):\n",
    "    return np.count_nonzero(board)\n",
    "\n",
    "def terminal_state(board):\n",
    "    if round_number(board) == MAX_ROUNDS:    # draw\n",
    "        return 0\n",
    "    if four_in_a_row(board, 1):\n",
    "        return 1\n",
    "    elif four_in_a_row(board, -1):\n",
    "        return -1\n",
    "    else:\n",
    "        return None"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## MCTS"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "from __future__ import annotations\n",
    "\n",
    "class Node:\n",
    "    def __init__(self, board: np.ndarray, player: int, parent: Node = None, move: int = None):\n",
    "        self.board = np.copy(board)\n",
    "        self.player = player    # player who did the previous move\n",
    "        self.parent = parent\n",
    "        self.move = move        # previous move that brought in this state\n",
    "        self.num_visits = 0\n",
    "        self.num_wins = 0\n",
    "        self.children = []\n",
    "        self.next_moves = valid_moves(board)\n",
    "\n",
    "    def selection(self):\n",
    "        def UCB1(node):\n",
    "            c = np.sqrt(2)\n",
    "            exploitation = node.num_wins / node.num_visits\n",
    "            exploration = c * np.sqrt(np.log(node.parent.num_visits) / node.num_visits)\n",
    "            return exploitation + exploration\n",
    "        \n",
    "        return max(self.children, key=UCB1)\n",
    "\n",
    "    def expand(self, move):\n",
    "        player = -self.player     \n",
    "        new_board = np.copy(self.board)\n",
    "        play(new_board, move, player)\n",
    "        self.next_moves.remove(move)\n",
    "        child = Node(new_board, player, self, move)\n",
    "        self.children.append(child)\n",
    "        return child\n",
    "\n",
    "    def simulate(self):\n",
    "        p = -self.player\n",
    "        board = np.copy(self.board)\n",
    "        while valid_moves(board):\n",
    "            move = np.random.choice(valid_moves(board))\n",
    "            play(board, move, p)\n",
    "            if four_in_a_row(board, p):\n",
    "                return p\n",
    "            p = -p\n",
    "        \n",
    "        return 0 # DRAW\n",
    "\n",
    "    def backpropagate(self, winner):\n",
    "        node = self\n",
    "        while node is not None:\n",
    "            if winner == 0:   # draw\n",
    "                node.num_wins += 0.5\n",
    "            elif winner == node.player:\n",
    "                node.num_wins += 1      \n",
    "            node.num_visits += 1\n",
    "            node = node.parent"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "def MCTS(board: np.ndarray, player: int, num_iterations: int = NUM_ITERATIONS):\n",
    "    # the player in the node is the one who did the previous move\n",
    "    root = Node(board, -player, parent=None, move=None)\n",
    "\n",
    "    for _ in range(num_iterations):\n",
    "        node = root\n",
    "\n",
    "        # SELECTION (tree traversal)\n",
    "        while len(node.children) != 0 and len(node.next_moves) == 0:  # until terminal or not fully expanded node\n",
    "            node = node.selection()\n",
    "\n",
    "        # EXPANSION        \n",
    "        if len(node.next_moves) > 0:\n",
    "            move = np.random.choice(node.next_moves)\n",
    "            node = node.expand(move)\n",
    "\n",
    "        # SIMULATION (ROLLOUT)\n",
    "        winner = terminal_state(node.board)\n",
    "        if winner is None:\n",
    "            winner = node.simulate()\n",
    "\n",
    "        # BACKPROPAGATION\n",
    "        node.backpropagate(winner)\n",
    "            \n",
    "    # Return most promising move from root (highest score)\n",
    "    best_node = max(root.children, key=lambda x: x.num_wins/x.num_visits)\n",
    "    return best_node.move"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Choosing Move"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "def choose_move(board: np.ndarray, player: int):\n",
    "    # FIRST/SECOND MOVE -> always play central\n",
    "    if not board.any() or round_number(board) == 1:\n",
    "        play(board, 3, player)\n",
    "        return 3\n",
    "\n",
    "    # THIRD MOVE -> always play in one of the 3 cells in the center\n",
    "    if round_number(board) == 2:\n",
    "        move = np.random.choice([2, 3, 4])\n",
    "        play(board, move, player)\n",
    "        return move\n",
    "\n",
    "    # MCTS \n",
    "    move = MCTS(board, player, NUM_ITERATIONS)\n",
    "\n",
    "    play(board, move, player)\n",
    "\n",
    "    return move"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Simulation examples"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### AI vs AI"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "A TURN -> 4\n",
      "- - - - - - - \n",
      "- - - - - - - \n",
      "- - - - - - - \n",
      "- - - - - - - \n",
      "- - - - - - - \n",
      "- - - A - - - \n",
      "\n",
      "B TURN -> 4\n",
      "- - - - - - - \n",
      "- - - - - - - \n",
      "- - - - - - - \n",
      "- - - - - - - \n",
      "- - - B - - - \n",
      "- - - A - - - \n",
      "\n",
      "A TURN -> 5\n",
      "- - - - - - - \n",
      "- - - - - - - \n",
      "- - - - - - - \n",
      "- - - - - - - \n",
      "- - - B - - - \n",
      "- - - A A - - \n",
      "\n",
      "B TURN -> 6\n",
      "- - - - - - - \n",
      "- - - - - - - \n",
      "- - - - - - - \n",
      "- - - - - - - \n",
      "- - - B - - - \n",
      "- - - A A B - \n",
      "\n",
      "A TURN -> 3\n",
      "- - - - - - - \n",
      "- - - - - - - \n",
      "- - - - - - - \n",
      "- - - - - - - \n",
      "- - - B - - - \n",
      "- - A A A B - \n",
      "\n",
      "B TURN -> 2\n",
      "- - - - - - - \n",
      "- - - - - - - \n",
      "- - - - - - - \n",
      "- - - - - - - \n",
      "- - - B - - - \n",
      "- B A A A B - \n",
      "\n",
      "A TURN -> 5\n",
      "- - - - - - - \n",
      "- - - - - - - \n",
      "- - - - - - - \n",
      "- - - - - - - \n",
      "- - - B A - - \n",
      "- B A A A B - \n",
      "\n",
      "B TURN -> 5\n",
      "- - - - - - - \n",
      "- - - - - - - \n",
      "- - - - - - - \n",
      "- - - - B - - \n",
      "- - - B A - - \n",
      "- B A A A B - \n",
      "\n",
      "A TURN -> 5\n",
      "- - - - - - - \n",
      "- - - - - - - \n",
      "- - - - A - - \n",
      "- - - - B - - \n",
      "- - - B A - - \n",
      "- B A A A B - \n",
      "\n",
      "B TURN -> 7\n",
      "- - - - - - - \n",
      "- - - - - - - \n",
      "- - - - A - - \n",
      "- - - - B - - \n",
      "- - - B A - - \n",
      "- B A A A B B \n",
      "\n",
      "A TURN -> 6\n",
      "- - - - - - - \n",
      "- - - - - - - \n",
      "- - - - A - - \n",
      "- - - - B - - \n",
      "- - - B A A - \n",
      "- B A A A B B \n",
      "\n",
      "B TURN -> 6\n",
      "- - - - - - - \n",
      "- - - - - - - \n",
      "- - - - A - - \n",
      "- - - - B B - \n",
      "- - - B A A - \n",
      "- B A A A B B \n",
      "\n",
      "A TURN -> 4\n",
      "- - - - - - - \n",
      "- - - - - - - \n",
      "- - - - A - - \n",
      "- - - A B B - \n",
      "- - - B A A - \n",
      "- B A A A B B \n",
      "\n",
      "B TURN -> 2\n",
      "- - - - - - - \n",
      "- - - - - - - \n",
      "- - - - A - - \n",
      "- - - A B B - \n",
      "- B - B A A - \n",
      "- B A A A B B \n",
      "\n",
      "A TURN -> 3\n",
      "- - - - - - - \n",
      "- - - - - - - \n",
      "- - - - A - - \n",
      "- - - A B B - \n",
      "- B A B A A - \n",
      "- B A A A B B \n",
      "\n",
      "B TURN -> 2\n",
      "- - - - - - - \n",
      "- - - - - - - \n",
      "- - - - A - - \n",
      "- B - A B B - \n",
      "- B A B A A - \n",
      "- B A A A B B \n",
      "\n",
      "A TURN -> 2\n",
      "- - - - - - - \n",
      "- - - - - - - \n",
      "- A - - A - - \n",
      "- B - A B B - \n",
      "- B A B A A - \n",
      "- B A A A B B \n",
      "\n",
      "B TURN -> 4\n",
      "- - - - - - - \n",
      "- - - - - - - \n",
      "- A - B A - - \n",
      "- B - A B B - \n",
      "- B A B A A - \n",
      "- B A A A B B \n",
      "\n",
      "A TURN -> 3\n",
      "- - - - - - - \n",
      "- - - - - - - \n",
      "- A - B A - - \n",
      "- B A A B B - \n",
      "- B A B A A - \n",
      "- B A A A B B \n",
      "\n",
      "B TURN -> 3\n",
      "- - - - - - - \n",
      "- - - - - - - \n",
      "- A B B A - - \n",
      "- B A A B B - \n",
      "- B A B A A - \n",
      "- B A A A B B \n",
      "\n",
      "A TURN -> 4\n",
      "- - - - - - - \n",
      "- - - A - - - \n",
      "- A B B A - - \n",
      "- B A A B B - \n",
      "- B A B A A - \n",
      "- B A A A B B \n",
      "\n",
      "B TURN -> 3\n",
      "- - - - - - - \n",
      "- - B A - - - \n",
      "- A B B A - - \n",
      "- B A A B B - \n",
      "- B A B A A - \n",
      "- B A A A B B \n",
      "\n",
      "A TURN -> 6\n",
      "- - - - - - - \n",
      "- - B A - - - \n",
      "- A B B A A - \n",
      "- B A A B B - \n",
      "- B A B A A - \n",
      "- B A A A B B \n",
      "\n",
      "B TURN -> 6\n",
      "- - - - - - - \n",
      "- - B A - B - \n",
      "- A B B A A - \n",
      "- B A A B B - \n",
      "- B A B A A - \n",
      "- B A A A B B \n",
      "\n",
      "A TURN -> 4\n",
      "- - - A - - - \n",
      "- - B A - B - \n",
      "- A B B A A - \n",
      "- B A A B B - \n",
      "- B A B A A - \n",
      "- B A A A B B \n",
      "\n",
      "B TURN -> 5\n",
      "- - - A - - - \n",
      "- - B A B B - \n",
      "- A B B A A - \n",
      "- B A A B B - \n",
      "- B A B A A - \n",
      "- B A A A B B \n",
      "\n",
      "A TURN -> 5\n",
      "- - - A A - - \n",
      "- - B A B B - \n",
      "- A B B A A - \n",
      "- B A A B B - \n",
      "- B A B A A - \n",
      "- B A A A B B \n",
      "\n",
      "B TURN -> 3\n",
      "- - B A A - - \n",
      "- - B A B B - \n",
      "- A B B A A - \n",
      "- B A A B B - \n",
      "- B A B A A - \n",
      "- B A A A B B \n",
      "\n",
      "A TURN -> 6\n",
      "- - B A A A - \n",
      "- - B A B B - \n",
      "- A B B A A - \n",
      "- B A A B B - \n",
      "- B A B A A - \n",
      "- B A A A B B \n",
      "\n",
      "B TURN -> 1\n",
      "- - B A A A - \n",
      "- - B A B B - \n",
      "- A B B A A - \n",
      "- B A A B B - \n",
      "- B A B A A - \n",
      "B B A A A B B \n",
      "\n",
      "A TURN -> 7\n",
      "- - B A A A - \n",
      "- - B A B B - \n",
      "- A B B A A - \n",
      "- B A A B B - \n",
      "- B A B A A A \n",
      "B B A A A B B \n",
      "\n",
      "B TURN -> 1\n",
      "- - B A A A - \n",
      "- - B A B B - \n",
      "- A B B A A - \n",
      "- B A A B B - \n",
      "B B A B A A A \n",
      "B B A A A B B \n",
      "\n",
      "A TURN -> 7\n",
      "- - B A A A - \n",
      "- - B A B B - \n",
      "- A B B A A - \n",
      "- B A A B B A \n",
      "B B A B A A A \n",
      "B B A A A B B \n",
      "\n",
      "B TURN -> 7\n",
      "- - B A A A - \n",
      "- - B A B B - \n",
      "- A B B A A B \n",
      "- B A A B B A \n",
      "B B A B A A A \n",
      "B B A A A B B \n",
      "\n",
      "A TURN -> 1\n",
      "- - B A A A - \n",
      "- - B A B B - \n",
      "- A B B A A B \n",
      "A B A A B B A \n",
      "B B A B A A A \n",
      "B B A A A B B \n",
      "\n",
      "B TURN -> 1\n",
      "- - B A A A - \n",
      "- - B A B B - \n",
      "B A B B A A B \n",
      "A B A A B B A \n",
      "B B A B A A A \n",
      "B B A A A B B \n",
      "\n",
      "A TURN -> 7\n",
      "- - B A A A - \n",
      "- - B A B B A \n",
      "B A B B A A B \n",
      "A B A A B B A \n",
      "B B A B A A A \n",
      "B B A A A B B \n",
      "\n",
      "B TURN -> 7\n",
      "- - B A A A B \n",
      "- - B A B B A \n",
      "B A B B A A B \n",
      "A B A A B B A \n",
      "B B A B A A A \n",
      "B B A A A B B \n",
      "\n",
      "A TURN -> 1\n",
      "- - B A A A B \n",
      "A - B A B B A \n",
      "B A B B A A B \n",
      "A B A A B B A \n",
      "B B A B A A A \n",
      "B B A A A B B \n",
      "\n",
      "B TURN -> 2\n",
      "- - B A A A B \n",
      "A B B A B B A \n",
      "B A B B A A B \n",
      "A B A A B B A \n",
      "B B A B A A A \n",
      "B B A A A B B \n",
      "\n",
      "A TURN -> 2\n",
      "- A B A A A B \n",
      "A B B A B B A \n",
      "B A B B A A B \n",
      "A B A A B B A \n",
      "B B A B A A A \n",
      "B B A A A B B \n",
      "\n",
      "B TURN -> 1\n",
      "B A B A A A B \n",
      "A B B A B B A \n",
      "B A B B A A B \n",
      "A B A A B B A \n",
      "B B A B A A A \n",
      "B B A A A B B \n",
      "\n",
      "DRAW\n"
     ]
    }
   ],
   "source": [
    "def main_AI_vs_AI():\n",
    "    board = initialize_board()\n",
    "    player = 1\n",
    "\n",
    "    while True:\n",
    "        move = choose_move(board, player)\n",
    "\n",
    "        print(f\"{PLAYERS[player]} TURN -> {move + 1}\")\n",
    "        display(board)\n",
    "        print()\n",
    "\n",
    "        if four_in_a_row(board, player):\n",
    "            print(f\"Player {PLAYERS[player]} WON\")\n",
    "            return\n",
    "\n",
    "        if round_number(board) == MAX_ROUNDS:\n",
    "            print(\"DRAW\")\n",
    "            return\n",
    "\n",
    "        player = -player\n",
    "\n",
    "\n",
    "main_AI_vs_AI()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Human vs AI"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def main_human_vs_AI():\n",
    "    board = initialize_board()\n",
    "    player = np.random.choice([-1, 1])\n",
    "\n",
    "    while True:\n",
    "        if player == 1:\n",
    "            move = choose_move(board, player)\n",
    "        else:\n",
    "            move = int(input(\"Enter column number (1-7): \"))\n",
    "            play(board, move-1, player)\n",
    "\n",
    "        print(f\"{PLAYERS[player]} TURN -> {move + 1}\")\n",
    "        display(board)\n",
    "        print()\n",
    "\n",
    "        if four_in_a_row(board, player):\n",
    "            print(f\"\\nPlayer {PLAYERS[player]} WON\")\n",
    "            return\n",
    "\n",
    "        if round_number(board) == MAX_ROUNDS:\n",
    "            print(\"DRAW\")\n",
    "            return\n",
    "\n",
    "        player = -player\n",
    "\n",
    "\n",
    "main_human_vs_AI()"
   ]
  }
 ],
 "metadata": {
  "interpreter": {
   "hash": "371627151b0642f6c185c280fa16312776bda89bbb063ce0f2d8135c657194f4"
  },
  "kernelspec": {
   "display_name": "Python 3.8.2 64-bit ('ci2021': conda)",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.7"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}