Connect4-AI.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Import Libraries and Define Game Environment"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Cell 1: Import Libraries and Define Game Environment\n",
    "\n",
    "import os\n",
    "import numpy as np\n",
    "import copy\n",
    "import math\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import torch.optim as optim\n",
    "from collections import deque\n",
    "from multiprocessing import Pool, cpu_count\n",
    "import random\n",
    "import logging\n",
    "import threading\n",
    "\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "logger = logging.getLogger()\n",
    "logger.setLevel(logging.INFO)\n",
    "# Create console handler and set level to info\n",
    "ch = logging.StreamHandler()\n",
    "ch.setLevel(logging.INFO)\n",
    "# Create formatter\n",
    "formatter = logging.Formatter('%(asctime)s - %(levelname)s - %(message)s')\n",
    "ch.setFormatter(formatter)\n",
    "# Add handler to logger\n",
    "logger.addHandler(ch)\n",
    "\n",
    "# For Logging and Visualization\n",
    "from torch.utils.tensorboard import SummaryWriter\n",
    "\n",
    "# For Interactive Widgets\n",
    "import ipywidgets as widgets\n",
    "from IPython.display import display, clear_output\n",
    "\n",
    "# Define the Connect Four game environment\n",
    "class ConnectFour:\n",
    "    ROWS = 6\n",
    "    COLS = 7\n",
    "\n",
    "    def __init__(self):\n",
    "        self.board = np.zeros((self.ROWS, self.COLS), dtype=int)\n",
    "        self.current_player = 1  # Player 1 starts\n",
    "\n",
    "    def make_move(self, col):\n",
    "        for row in reversed(range(self.ROWS)):\n",
    "            if self.board[row, col] == 0:\n",
    "                self.board[row, col] = self.current_player\n",
    "                self.current_player *= -1  # Switch player\n",
    "                return True\n",
    "        return False  # Column is full\n",
    "\n",
    "    def valid_moves(self):\n",
    "        return [col for col in range(self.COLS) if self.board[0, col] == 0]\n",
    "\n",
    "    def is_full(self):\n",
    "        return np.all(self.board != 0)\n",
    "\n",
    "    def check_winner(self):\n",
    "        # Check horizontal locations for win\n",
    "        for c in range(self.COLS - 3):\n",
    "            for r in range(self.ROWS):\n",
    "                piece = self.board[r][c]\n",
    "                if piece != 0 and all(self.board[r][c + i] == piece for i in range(4)):\n",
    "                    return piece\n",
    "\n",
    "        # Check vertical locations for win\n",
    "        for c in range(self.COLS):\n",
    "            for r in range(self.ROWS - 3):\n",
    "                piece = self.board[r][c]\n",
    "                if piece != 0 and all(self.board[r + i][c] == piece for i in range(4)):\n",
    "                    return piece\n",
    "\n",
    "        # Check positively sloped diagonals\n",
    "        for c in range(self.COLS - 3):\n",
    "            for r in range(self.ROWS - 3):\n",
    "                piece = self.board[r][c]\n",
    "                if piece != 0 and all(self.board[r + i][c + i] == piece for i in range(4)):\n",
    "                    return piece\n",
    "\n",
    "        # Check negatively sloped diagonals\n",
    "        for c in range(self.COLS - 3):\n",
    "            for r in range(3, self.ROWS):\n",
    "                piece = self.board[r][c]\n",
    "                if piece != 0 and all(self.board[r - i][c + i] == piece for i in range(4)):\n",
    "                    return piece\n",
    "\n",
    "        return 0  # No winner\n",
    "\n",
    "    def reset(self):\n",
    "        self.board = np.zeros((self.ROWS, self.COLS), dtype=int)\n",
    "        self.current_player = 1"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Define Neural Network with Residual Blocks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Cell 2: Define Neural Network with Residual Blocks\n",
    "\n",
    "# Define a Residual Block\n",
    "class ResidualBlock(nn.Module):\n",
    "    def __init__(self, channels):\n",
    "        super(ResidualBlock, self).__init__()\n",
    "        self.conv1 = nn.Conv2d(channels, channels, kernel_size=3, padding=1)\n",
    "        self.bn1 = nn.BatchNorm2d(channels)\n",
    "        self.conv2 = nn.Conv2d(channels, channels, kernel_size=3, padding=1)\n",
    "        self.bn2 = nn.BatchNorm2d(channels)\n",
    "\n",
    "    def forward(self, x):\n",
    "        residual = x\n",
    "        out = F.relu(self.bn1(self.conv1(x)))\n",
    "        out = self.bn2(self.conv2(out))\n",
    "        out += residual  # Out-of-place addition to prevent in-place errors\n",
    "        out = F.relu(out)\n",
    "        return out\n",
    "\n",
    "# Define the ConnectFour Neural Network with Residual Blocks\n",
    "class ConnectNet(nn.Module):\n",
    "    def __init__(self, num_residual_blocks=6):\n",
    "        super(ConnectNet, self).__init__()\n",
    "        self.conv1 = nn.Conv2d(1, 64, kernel_size=3, padding=1)\n",
    "        self.bn1 = nn.BatchNorm2d(64)\n",
    "        # Stack multiple Residual Blocks\n",
    "        self.residual_blocks = nn.Sequential(\n",
    "            *[ResidualBlock(64) for _ in range(num_residual_blocks)]\n",
    "        )\n",
    "        self.conv_final = nn.Conv2d(64, 64, kernel_size=3, padding=1)\n",
    "        self.bn_final = nn.BatchNorm2d(64)\n",
    "        self.fc1 = nn.Linear(64 * 6 * 7, 512)\n",
    "        self.dropout = nn.Dropout(p=0.6)  # Increased dropout for regularization\n",
    "        self.fc_policy = nn.Linear(512, 7)\n",
    "        self.fc_value = nn.Linear(512, 1)\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = x.view(-1, 1, 6, 7)\n",
    "        x = F.relu(self.bn1(self.conv1(x)))\n",
    "        x = self.residual_blocks(x)\n",
    "        x = F.relu(self.bn_final(self.conv_final(x)))\n",
    "        x = x.view(-1, 64 * 6 * 7)\n",
    "        x = self.dropout(F.relu(self.fc1(x)))\n",
    "        policy_logits = self.fc_policy(x)\n",
    "        value = torch.tanh(self.fc_value(x))\n",
    "        return policy_logits, value"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Define MCTS and Related Functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Cell 3: Define MCTS and Related Functions\n",
    "\n",
    "# Define the MCTS Node\n",
    "class MCTSNode:\n",
    "    def __init__(self, state, parent=None):\n",
    "        self.state = copy.deepcopy(state)\n",
    "        self.parent = parent\n",
    "        self.children = {}\n",
    "        self.visits = 0\n",
    "        self.value = 0.0\n",
    "        self.prior = 0.0\n",
    "\n",
    "# UCB Score Calculation\n",
    "def ucb_score(parent, child, c_puct=2.0):\n",
    "    prior_score = c_puct * child.prior * math.sqrt(parent.visits) / (1 + child.visits)\n",
    "    value_score = child.value / (1 + child.visits)\n",
    "    return prior_score - value_score\n",
    "\n",
    "# Backpropagate the value up the path\n",
    "def backpropagate(path, value):\n",
    "    for node in reversed(path):\n",
    "        node.visits += 1\n",
    "        node.value += value\n",
    "        value = -value  # Flip the value for the opponent\n",
    "\n",
    "# Monte Carlo Tree Search\n",
    "def mcts_search(root, net, num_simulations=800):\n",
    "    for _ in range(num_simulations):\n",
    "        node = root\n",
    "        path = []\n",
    "\n",
    "        # Selection\n",
    "        while node.children:\n",
    "            # Select the child with the highest UCB score\n",
    "            child = max(node.children.values(), key=lambda child: ucb_score(node, child))\n",
    "            node = child\n",
    "            path.append(node)\n",
    "\n",
    "        # Expansion\n",
    "        winner = node.state.check_winner()\n",
    "        if winner == 0 and not node.state.is_full():\n",
    "            # Prepare the state tensor\n",
    "            state_tensor = torch.tensor(node.state.board, dtype=torch.float32).unsqueeze(0).to(device)\n",
    "            net.eval()\n",
    "            with torch.no_grad():\n",
    "                policy_logits, value = net(state_tensor)\n",
    "            net.train()\n",
    "            policy = F.softmax(policy_logits, dim=1).cpu().numpy()[0]\n",
    "\n",
    "            valid_moves = node.state.valid_moves()\n",
    "            policy = {idx: policy[idx] for idx in valid_moves}\n",
    "            policy_sum = sum(policy.values())\n",
    "            for idx in policy:\n",
    "                policy[idx] /= policy_sum  # Normalize the probabilities\n",
    "\n",
    "            # Add Dirichlet noise at the root node for exploration\n",
    "            if node == root:\n",
    "                dirichlet_alpha = 0.3\n",
    "                epsilon = 0.25\n",
    "                dirichlet_noise = np.random.dirichlet([dirichlet_alpha] * len(valid_moves))\n",
    "                for i, idx in enumerate(valid_moves):\n",
    "                    policy[idx] = (1 - epsilon) * policy[idx] + epsilon * dirichlet_noise[i]\n",
    "\n",
    "            # Expand children\n",
    "            for idx in valid_moves:\n",
    "                child_state = copy.deepcopy(node.state)\n",
    "                child_state.make_move(idx)\n",
    "                child_node = MCTSNode(child_state, node)\n",
    "                child_node.prior = policy[idx]\n",
    "                node.children[idx] = child_node\n",
    "            # Use the value estimate from the neural network\n",
    "            backpropagate(path + [node], value.item())\n",
    "        else:\n",
    "            # Terminal node\n",
    "            value = winner if winner != 0 else 0\n",
    "            backpropagate(path + [node], value)\n",
    "    # Choose the move with the most visits\n",
    "    best_move = max(root.children.items(), key=lambda item: item[1].visits)[0]\n",
    "    return best_move"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Define Replay Buffer, Data Augmentation, and Self-Play Functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Cell 4: Define Replay Buffer, Data Augmentation, and Self-Play Functions\n",
    "\n",
    "# Experience Replay Buffer\n",
    "class ReplayBuffer:\n",
    "    def __init__(self, max_size=1000000):\n",
    "        self.buffer = deque(maxlen=max_size)\n",
    "    \n",
    "    def add(self, memory):\n",
    "        self.buffer.extend(memory)\n",
    "    \n",
    "    def sample(self, batch_size):\n",
    "        if len(self.buffer) < batch_size:\n",
    "            return list(self.buffer)\n",
    "        else:\n",
    "            return random.sample(self.buffer, batch_size)\n",
    "    \n",
    "    def __len__(self):\n",
    "        return len(self.buffer)\n",
    "\n",
    "# Data Augmentation Functions\n",
    "\n",
    "def augment_board(board):\n",
    "    \"\"\"\n",
    "    Generate augmented versions of the board by applying horizontal reflections.\n",
    "    Returns a list of augmented boards.\n",
    "    \"\"\"\n",
    "    augmented_boards = []\n",
    "    # Original board\n",
    "    augmented_boards.append(board)\n",
    "    # Horizontal flip\n",
    "    flipped = np.fliplr(board)\n",
    "    augmented_boards.append(flipped)\n",
    "    return augmented_boards\n",
    "\n",
    "def augment_memory(memory):\n",
    "    \"\"\"\n",
    "    Apply data augmentation to the entire memory.\n",
    "    Returns a new augmented memory list.\n",
    "    \"\"\"\n",
    "    augmented_memory = []\n",
    "    for state, probs, reward in memory:\n",
    "        augmented_states = augment_board(state)\n",
    "        for aug_state in augmented_states:\n",
    "            augmented_memory.append((aug_state, probs, reward))\n",
    "    return augmented_memory\n",
    "\n",
    "# Self-Play Function\n",
    "def self_play(net, num_games=50, num_simulations=800):\n",
    "    memory = []\n",
    "    win_cnt = {1: 0, -1: 0, 'draw': 0}\n",
    "    for game_num in range(num_games):\n",
    "        game = ConnectFour()\n",
    "        # Randomize starting player\n",
    "        game.current_player = np.random.choice([1, -1])\n",
    "        states = []\n",
    "        mcts_probs = []\n",
    "        current_players = []\n",
    "        while True:\n",
    "            root = MCTSNode(game)\n",
    "            move = mcts_search(root, net, num_simulations=num_simulations)\n",
    "            states.append(game.board.copy())\n",
    "            mcts_prob = np.zeros(7)\n",
    "            for idx, child in root.children.items():\n",
    "                mcts_prob[idx] = child.visits\n",
    "            if np.sum(mcts_prob) > 0:\n",
    "                mcts_prob = mcts_prob / np.sum(mcts_prob)\n",
    "            else:\n",
    "                # If no children were expanded, use uniform probabilities\n",
    "                mcts_prob = np.ones(7) / 7\n",
    "            mcts_probs.append(mcts_prob)\n",
    "            current_players.append(game.current_player)\n",
    "            valid_move = game.make_move(move)\n",
    "            if not valid_move:\n",
    "                logger.warning(f\"Invalid move attempted at column {move} by player {game.current_player}.\")\n",
    "                break  # Avoid infinite loops\n",
    "            winner = game.check_winner()\n",
    "            if winner != 0 or game.is_full():\n",
    "                break\n",
    "        # Assign values to the game states\n",
    "        if winner == 0:\n",
    "            win_cnt['draw'] += 1\n",
    "        else:\n",
    "            win_cnt[winner] += 1\n",
    "        for i in range(len(states)):\n",
    "            reward = winner if winner == current_players[i] else -winner\n",
    "            memory.append((states[i], mcts_probs[i], reward))\n",
    "    return memory, win_cnt\n",
    "\n",
    "# Training Function\n",
    "def train(net, replay_buffer, optimizer, batch_size=64, epochs=1, clip_grad=1.0):\n",
    "    net.train()\n",
    "    dataset_size = len(replay_buffer.buffer)\n",
    "    if dataset_size < batch_size:\n",
    "        logger.info(\"Not enough data to train. Skipping this epoch.\")\n",
    "        return 0, 0, 0  # Not enough data to train\n",
    "    total_value_loss = 0.0\n",
    "    total_policy_loss = 0.0\n",
    "    total_loss = 0.0\n",
    "    num_batches = 0\n",
    "\n",
    "    for epoch in range(epochs):\n",
    "        # Shuffle the data\n",
    "        perm = np.random.permutation(dataset_size)\n",
    "        for i in range(0, dataset_size, batch_size):\n",
    "            end_idx = min(i + batch_size, dataset_size)\n",
    "            batch_indices = perm[i:end_idx]\n",
    "            batch = [replay_buffer.buffer[idx] for idx in batch_indices]\n",
    "            batch_states, batch_probs, batch_rewards = zip(*batch)\n",
    "            \n",
    "            # Validate and filter states\n",
    "            valid_batch = []\n",
    "            for state, prob, reward in zip(batch_states, batch_probs, batch_rewards):\n",
    "                state_array = np.array(state)\n",
    "                if state_array.shape == (6, 7):\n",
    "                    valid_batch.append((state_array, prob, reward))\n",
    "                else:\n",
    "                    logger.warning(f\"Invalid state shape detected: {state_array.shape}. Skipping this state.\")\n",
    "            \n",
    "            if not valid_batch:\n",
    "                logger.warning(\"No valid game states in this batch. Skipping training for this epoch.\")\n",
    "                continue\n",
    "            \n",
    "            # Unzip the valid batch\n",
    "            valid_batch_states, valid_batch_probs, valid_batch_rewards = zip(*valid_batch)\n",
    "            \n",
    "            # Convert to tensors\n",
    "            state_tensor = torch.tensor(np.array(valid_batch_states), dtype=torch.float32).to(device)\n",
    "            mcts_prob_tensor = torch.tensor(np.array(valid_batch_probs), dtype=torch.float32).to(device)\n",
    "            reward_tensor = torch.tensor(np.array(valid_batch_rewards), dtype=torch.float32).to(device)\n",
    "            \n",
    "            # Forward pass\n",
    "            policy_logits, value = net(state_tensor)\n",
    "            \n",
    "            # Calculate losses\n",
    "            value_loss = F.mse_loss(value.squeeze(), reward_tensor)\n",
    "            policy_loss = -torch.mean(torch.sum(mcts_prob_tensor * F.log_softmax(policy_logits, dim=1), dim=1))\n",
    "            loss = value_loss + policy_loss\n",
    "            \n",
    "            # Backward pass and optimization\n",
    "            optimizer.zero_grad()\n",
    "            loss.backward()\n",
    "            torch.nn.utils.clip_grad_norm_(net.parameters(), max_norm=clip_grad)\n",
    "            optimizer.step()\n",
    "            \n",
    "            total_value_loss += value_loss.item()\n",
    "            total_policy_loss += policy_loss.item()\n",
    "            total_loss += loss.item()\n",
    "            num_batches += 1\n",
    "    \n",
    "    if num_batches > 0:\n",
    "        avg_value_loss = total_value_loss / num_batches\n",
    "        avg_policy_loss = total_policy_loss / num_batches\n",
    "        avg_loss = total_loss / num_batches\n",
    "    else:\n",
    "        avg_value_loss = 0\n",
    "        avg_policy_loss = 0\n",
    "        avg_loss = 0\n",
    "    return avg_value_loss, avg_policy_loss, avg_loss"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Define Evaluation and Checkpointing Functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Cell 5: Define Evaluation and Checkpointing Functions\n",
    "\n",
    "# Evaluation Function\n",
    "def evaluate(net, num_games=20, num_simulations=800):\n",
    "    net.eval()\n",
    "    win_cnt = {1: 0, -1: 0, 'draw': 0}\n",
    "    for game_num in range(num_games):\n",
    "        game = ConnectFour()\n",
    "        game.current_player = np.random.choice([1, -1])  # Randomize starting player\n",
    "        while True:\n",
    "            if game.current_player == 1:\n",
    "                # AI move\n",
    "                root = MCTSNode(game)\n",
    "                ai_move = mcts_search(root, net, num_simulations=num_simulations)\n",
    "                valid_move = game.make_move(ai_move)\n",
    "                if not valid_move:\n",
    "                    logger.warning(f\"AI attempted invalid move at column {ai_move}.\")\n",
    "                    break\n",
    "            else:\n",
    "                # Opponent move (heuristic-based)\n",
    "                valid_moves = game.valid_moves()\n",
    "                # Simple heuristic: choose the column with the least filled spots\n",
    "                heights = [np.sum(game.board[:, col] != 0) for col in valid_moves]\n",
    "                min_height = min(heights)\n",
    "                candidate_cols = [col for col, height in zip(valid_moves, heights) if height == min_height]\n",
    "                move = np.random.choice(candidate_cols)\n",
    "                game.make_move(move)\n",
    "            winner = game.check_winner()\n",
    "            if winner != 0 or game.is_full():\n",
    "                break\n",
    "        if winner == 0:\n",
    "            win_cnt['draw'] += 1\n",
    "        else:\n",
    "            win_cnt[winner] += 1\n",
    "    return win_cnt\n",
    "\n",
    "# Save Checkpoint\n",
    "def save_checkpoint(state, filename='connect4_model.pth'):\n",
    "    torch.save(state, filename)\n",
    "    logger.info(f\"Checkpoint saved to {filename}\")\n",
    "\n",
    "# Load Checkpoint\n",
    "def load_checkpoint(filename='connect4_model.pth'):\n",
    "    if not os.path.exists(filename):\n",
    "        logger.warning(f\"No checkpoint found at {filename}.\")\n",
    "        return None\n",
    "    checkpoint = torch.load(filename, map_location=device)\n",
    "    logger.info(f\"Checkpoint loaded from {filename}\")\n",
    "    return checkpoint"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Logging and Visualisation Setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Reusing TensorBoard on port 6008 (pid 15508), started 1 day, 20:55:45 ago. (Use '!kill 15508' to kill it.)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "\n",
       "      <iframe id=\"tensorboard-frame-ab5aa5a9d10800ef\" width=\"100%\" height=\"800\" frameborder=\"0\">\n",
       "      </iframe>\n",
       "      <script>\n",
       "        (function() {\n",
       "          const frame = document.getElementById(\"tensorboard-frame-ab5aa5a9d10800ef\");\n",
       "          const url = new URL(\"http://localhost\");\n",
       "          const port = 6008;\n",
       "          if (port) {\n",
       "            url.port = port;\n",
       "          }\n",
       "          frame.src = url;\n",
       "        })();\n",
       "      </script>\n",
       "    "
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Cell 6: Logging and Visualization Setup\n",
    "\n",
    "# Initialize TensorBoard writer\n",
    "writer = SummaryWriter(log_dir='connect4_logs')\n",
    "\n",
    "# Function to log training metrics\n",
    "def log_training_metrics(iteration, avg_value_loss, avg_policy_loss, avg_loss, self_play_stats, eval_stats=None):\n",
    "    writer.add_scalar('Loss/Value', avg_value_loss, iteration)\n",
    "    writer.add_scalar('Loss/Policy', avg_policy_loss, iteration)\n",
    "    writer.add_scalar('Loss/Total', avg_loss, iteration)\n",
    "    \n",
    "    # Log self-play statistics\n",
    "    for key, value in self_play_stats.items():\n",
    "        writer.add_scalar(f'SelfPlay/{key}', value, iteration)\n",
    "    \n",
    "    # Log evaluation statistics if available\n",
    "    if eval_stats:\n",
    "        for key, value in eval_stats.items():\n",
    "            writer.add_scalar(f'Evaluation/{key}', value, iteration)\n",
    "    \n",
    "    writer.flush()  # Ensure all pending logs are written\n",
    "\n",
    "# Launch TensorBoard within Jupyter\n",
    "%load_ext tensorboard\n",
    "%tensorboard --logdir='C:\\Users\\Hephzibah\\OneDrive\\Documents\\VScodeprojects\\Connect4_project\\connect4_logs'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Hyperparameter Configuration"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Cell 7: Hyperparameter Configuration\n",
    "\n",
    "# Define Hyperparameters\n",
    "class Hyperparameters:\n",
    "    # Training Parameters\n",
    "    num_residual_blocks = 6\n",
    "    learning_rate = 0.001\n",
    "    weight_decay = 1e-4\n",
    "    batch_size = 64\n",
    "    epochs = 1\n",
    "    clip_grad = 1.0\n",
    "    total_iterations = 1000\n",
    "    \n",
    "    # Self-Play Parameters\n",
    "    num_self_play_games = 50\n",
    "    num_mcts_simulations = 800\n",
    "    \n",
    "    # Replay Buffer\n",
    "    replay_buffer_max_size = 1000000\n",
    "    \n",
    "    # Evaluation Parameters\n",
    "    evaluation_interval = 10\n",
    "    num_evaluation_games = 20\n",
    "    \n",
    "    # MCTS Parameters\n",
    "    c_puct = 2.0\n",
    "    dirichlet_alpha = 0.3\n",
    "    epsilon = 0.25\n",
    "\n",
    "# Instantiate Hyperparameters\n",
    "hp = Hyperparameters()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Training Loop"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "C:\\Users\\Hephzibah\\AppData\\Local\\Temp\\ipykernel_15336\\3456340580.py:47: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.\n",
      "  checkpoint = torch.load(filename, map_location=device)\n",
      "2024-11-22 13:04:12,736 - INFO - Checkpoint loaded from connect4_model.pth\n",
      "2024-11-22 13:04:12,745 - INFO - Loaded model from iteration 33\n",
      "2024-11-22 13:04:12,745 - INFO - \n",
      "=== Iteration 34/1000 ===\n"
     ]
    },
    {
     "ename": "KeyboardInterrupt",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[1;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
      "Cell \u001b[1;32mIn[8], line 47\u001b[0m\n\u001b[0;32m     45\u001b[0m logger\u001b[38;5;241m.\u001b[39minfo(\u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;130;01m\\n\u001b[39;00m\u001b[38;5;124m=== Iteration \u001b[39m\u001b[38;5;132;01m{\u001b[39;00miteration\u001b[38;5;250m \u001b[39m\u001b[38;5;241m+\u001b[39m\u001b[38;5;250m \u001b[39m\u001b[38;5;241m1\u001b[39m\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m/\u001b[39m\u001b[38;5;132;01m{\u001b[39;00mtotal_iterations\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m ===\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n\u001b[0;32m     46\u001b[0m \u001b[38;5;66;03m# Self-Play\u001b[39;00m\n\u001b[1;32m---> 47\u001b[0m memory, self_play_stats \u001b[38;5;241m=\u001b[39m \u001b[43mself_play\u001b[49m\u001b[43m(\u001b[49m\u001b[43mnet\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mnum_games\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mhp\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mnum_self_play_games\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mnum_simulations\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mhp\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mnum_mcts_simulations\u001b[49m\u001b[43m)\u001b[49m\n\u001b[0;32m     49\u001b[0m \u001b[38;5;66;03m# Apply Data Augmentation\u001b[39;00m\n\u001b[0;32m     50\u001b[0m augmented_memory \u001b[38;5;241m=\u001b[39m augment_memory(memory)\n",
      "Cell \u001b[1;32mIn[4], line 60\u001b[0m, in \u001b[0;36mself_play\u001b[1;34m(net, num_games, num_simulations)\u001b[0m\n\u001b[0;32m     58\u001b[0m \u001b[38;5;28;01mwhile\u001b[39;00m \u001b[38;5;28;01mTrue\u001b[39;00m:\n\u001b[0;32m     59\u001b[0m     root \u001b[38;5;241m=\u001b[39m MCTSNode(game)\n\u001b[1;32m---> 60\u001b[0m     move \u001b[38;5;241m=\u001b[39m \u001b[43mmcts_search\u001b[49m\u001b[43m(\u001b[49m\u001b[43mroot\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mnet\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mnum_simulations\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mnum_simulations\u001b[49m\u001b[43m)\u001b[49m\n\u001b[0;32m     61\u001b[0m     states\u001b[38;5;241m.\u001b[39mappend(game\u001b[38;5;241m.\u001b[39mboard\u001b[38;5;241m.\u001b[39mcopy())\n\u001b[0;32m     62\u001b[0m     mcts_prob \u001b[38;5;241m=\u001b[39m np\u001b[38;5;241m.\u001b[39mzeros(\u001b[38;5;241m7\u001b[39m)\n",
      "Cell \u001b[1;32mIn[3], line 46\u001b[0m, in \u001b[0;36mmcts_search\u001b[1;34m(root, net, num_simulations)\u001b[0m\n\u001b[0;32m     44\u001b[0m net\u001b[38;5;241m.\u001b[39meval()\n\u001b[0;32m     45\u001b[0m \u001b[38;5;28;01mwith\u001b[39;00m torch\u001b[38;5;241m.\u001b[39mno_grad():\n\u001b[1;32m---> 46\u001b[0m     policy_logits, value \u001b[38;5;241m=\u001b[39m \u001b[43mnet\u001b[49m\u001b[43m(\u001b[49m\u001b[43mstate_tensor\u001b[49m\u001b[43m)\u001b[49m\n\u001b[0;32m     47\u001b[0m net\u001b[38;5;241m.\u001b[39mtrain()\n\u001b[0;32m     48\u001b[0m policy \u001b[38;5;241m=\u001b[39m F\u001b[38;5;241m.\u001b[39msoftmax(policy_logits, dim\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m1\u001b[39m)\u001b[38;5;241m.\u001b[39mcpu()\u001b[38;5;241m.\u001b[39mnumpy()[\u001b[38;5;241m0\u001b[39m]\n",
      "File \u001b[1;32mc:\\Users\\Hephzibah\\AppData\\Local\\Programs\\Python\\Python310\\lib\\site-packages\\torch\\nn\\modules\\module.py:1736\u001b[0m, in \u001b[0;36mModule._wrapped_call_impl\u001b[1;34m(self, *args, **kwargs)\u001b[0m\n\u001b[0;32m   1734\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_compiled_call_impl(\u001b[38;5;241m*\u001b[39margs, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs)  \u001b[38;5;66;03m# type: ignore[misc]\u001b[39;00m\n\u001b[0;32m   1735\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[1;32m-> 1736\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_call_impl(\u001b[38;5;241m*\u001b[39margs, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs)\n",
      "File \u001b[1;32mc:\\Users\\Hephzibah\\AppData\\Local\\Programs\\Python\\Python310\\lib\\site-packages\\torch\\nn\\modules\\module.py:1747\u001b[0m, in \u001b[0;36mModule._call_impl\u001b[1;34m(self, *args, **kwargs)\u001b[0m\n\u001b[0;32m   1742\u001b[0m \u001b[38;5;66;03m# If we don't have any hooks, we want to skip the rest of the logic in\u001b[39;00m\n\u001b[0;32m   1743\u001b[0m \u001b[38;5;66;03m# this function, and just call forward.\u001b[39;00m\n\u001b[0;32m   1744\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m (\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_backward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_backward_pre_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_pre_hooks\n\u001b[0;32m   1745\u001b[0m         \u001b[38;5;129;01mor\u001b[39;00m _global_backward_pre_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_backward_hooks\n\u001b[0;32m   1746\u001b[0m         \u001b[38;5;129;01mor\u001b[39;00m _global_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_forward_pre_hooks):\n\u001b[1;32m-> 1747\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m forward_call(\u001b[38;5;241m*\u001b[39margs, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs)\n\u001b[0;32m   1749\u001b[0m result \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m\n\u001b[0;32m   1750\u001b[0m called_always_called_hooks \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mset\u001b[39m()\n",
      "Cell \u001b[1;32mIn[2], line 40\u001b[0m, in \u001b[0;36mConnectNet.forward\u001b[1;34m(self, x)\u001b[0m\n\u001b[0;32m     38\u001b[0m x \u001b[38;5;241m=\u001b[39m x\u001b[38;5;241m.\u001b[39mview(\u001b[38;5;241m-\u001b[39m\u001b[38;5;241m1\u001b[39m, \u001b[38;5;241m1\u001b[39m, \u001b[38;5;241m6\u001b[39m, \u001b[38;5;241m7\u001b[39m)\n\u001b[0;32m     39\u001b[0m x \u001b[38;5;241m=\u001b[39m F\u001b[38;5;241m.\u001b[39mrelu(\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mbn1(\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mconv1(x)))\n\u001b[1;32m---> 40\u001b[0m x \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mresidual_blocks\u001b[49m\u001b[43m(\u001b[49m\u001b[43mx\u001b[49m\u001b[43m)\u001b[49m\n\u001b[0;32m     41\u001b[0m x \u001b[38;5;241m=\u001b[39m F\u001b[38;5;241m.\u001b[39mrelu(\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mbn_final(\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mconv_final(x)))\n\u001b[0;32m     42\u001b[0m x \u001b[38;5;241m=\u001b[39m x\u001b[38;5;241m.\u001b[39mview(\u001b[38;5;241m-\u001b[39m\u001b[38;5;241m1\u001b[39m, \u001b[38;5;241m64\u001b[39m \u001b[38;5;241m*\u001b[39m \u001b[38;5;241m6\u001b[39m \u001b[38;5;241m*\u001b[39m \u001b[38;5;241m7\u001b[39m)\n",
      "File \u001b[1;32mc:\\Users\\Hephzibah\\AppData\\Local\\Programs\\Python\\Python310\\lib\\site-packages\\torch\\nn\\modules\\module.py:1736\u001b[0m, in \u001b[0;36mModule._wrapped_call_impl\u001b[1;34m(self, *args, **kwargs)\u001b[0m\n\u001b[0;32m   1734\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_compiled_call_impl(\u001b[38;5;241m*\u001b[39margs, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs)  \u001b[38;5;66;03m# type: ignore[misc]\u001b[39;00m\n\u001b[0;32m   1735\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[1;32m-> 1736\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_call_impl(\u001b[38;5;241m*\u001b[39margs, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs)\n",
      "File \u001b[1;32mc:\\Users\\Hephzibah\\AppData\\Local\\Programs\\Python\\Python310\\lib\\site-packages\\torch\\nn\\modules\\module.py:1747\u001b[0m, in \u001b[0;36mModule._call_impl\u001b[1;34m(self, *args, **kwargs)\u001b[0m\n\u001b[0;32m   1742\u001b[0m \u001b[38;5;66;03m# If we don't have any hooks, we want to skip the rest of the logic in\u001b[39;00m\n\u001b[0;32m   1743\u001b[0m \u001b[38;5;66;03m# this function, and just call forward.\u001b[39;00m\n\u001b[0;32m   1744\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m (\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_backward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_backward_pre_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_pre_hooks\n\u001b[0;32m   1745\u001b[0m         \u001b[38;5;129;01mor\u001b[39;00m _global_backward_pre_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_backward_hooks\n\u001b[0;32m   1746\u001b[0m         \u001b[38;5;129;01mor\u001b[39;00m _global_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_forward_pre_hooks):\n\u001b[1;32m-> 1747\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m forward_call(\u001b[38;5;241m*\u001b[39margs, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs)\n\u001b[0;32m   1749\u001b[0m result \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m\n\u001b[0;32m   1750\u001b[0m called_always_called_hooks \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mset\u001b[39m()\n",
      "File \u001b[1;32mc:\\Users\\Hephzibah\\AppData\\Local\\Programs\\Python\\Python310\\lib\\site-packages\\torch\\nn\\modules\\container.py:250\u001b[0m, in \u001b[0;36mSequential.forward\u001b[1;34m(self, input)\u001b[0m\n\u001b[0;32m    248\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mforward\u001b[39m(\u001b[38;5;28mself\u001b[39m, \u001b[38;5;28minput\u001b[39m):\n\u001b[0;32m    249\u001b[0m     \u001b[38;5;28;01mfor\u001b[39;00m module \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28mself\u001b[39m:\n\u001b[1;32m--> 250\u001b[0m         \u001b[38;5;28minput\u001b[39m \u001b[38;5;241m=\u001b[39m \u001b[43mmodule\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43minput\u001b[39;49m\u001b[43m)\u001b[49m\n\u001b[0;32m    251\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28minput\u001b[39m\n",
      "File \u001b[1;32mc:\\Users\\Hephzibah\\AppData\\Local\\Programs\\Python\\Python310\\lib\\site-packages\\torch\\nn\\modules\\module.py:1736\u001b[0m, in \u001b[0;36mModule._wrapped_call_impl\u001b[1;34m(self, *args, **kwargs)\u001b[0m\n\u001b[0;32m   1734\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_compiled_call_impl(\u001b[38;5;241m*\u001b[39margs, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs)  \u001b[38;5;66;03m# type: ignore[misc]\u001b[39;00m\n\u001b[0;32m   1735\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[1;32m-> 1736\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_call_impl(\u001b[38;5;241m*\u001b[39margs, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs)\n",
      "File \u001b[1;32mc:\\Users\\Hephzibah\\AppData\\Local\\Programs\\Python\\Python310\\lib\\site-packages\\torch\\nn\\modules\\module.py:1747\u001b[0m, in \u001b[0;36mModule._call_impl\u001b[1;34m(self, *args, **kwargs)\u001b[0m\n\u001b[0;32m   1742\u001b[0m \u001b[38;5;66;03m# If we don't have any hooks, we want to skip the rest of the logic in\u001b[39;00m\n\u001b[0;32m   1743\u001b[0m \u001b[38;5;66;03m# this function, and just call forward.\u001b[39;00m\n\u001b[0;32m   1744\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m (\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_backward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_backward_pre_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_pre_hooks\n\u001b[0;32m   1745\u001b[0m         \u001b[38;5;129;01mor\u001b[39;00m _global_backward_pre_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_backward_hooks\n\u001b[0;32m   1746\u001b[0m         \u001b[38;5;129;01mor\u001b[39;00m _global_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_forward_pre_hooks):\n\u001b[1;32m-> 1747\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m forward_call(\u001b[38;5;241m*\u001b[39margs, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs)\n\u001b[0;32m   1749\u001b[0m result \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m\n\u001b[0;32m   1750\u001b[0m called_always_called_hooks \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mset\u001b[39m()\n",
      "Cell \u001b[1;32mIn[2], line 14\u001b[0m, in \u001b[0;36mResidualBlock.forward\u001b[1;34m(self, x)\u001b[0m\n\u001b[0;32m     12\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mforward\u001b[39m(\u001b[38;5;28mself\u001b[39m, x):\n\u001b[0;32m     13\u001b[0m     residual \u001b[38;5;241m=\u001b[39m x\n\u001b[1;32m---> 14\u001b[0m     out \u001b[38;5;241m=\u001b[39m F\u001b[38;5;241m.\u001b[39mrelu(\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mbn1\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mconv1\u001b[49m\u001b[43m(\u001b[49m\u001b[43mx\u001b[49m\u001b[43m)\u001b[49m\u001b[43m)\u001b[49m)\n\u001b[0;32m     15\u001b[0m     out \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mbn2(\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mconv2(out))\n\u001b[0;32m     16\u001b[0m     out \u001b[38;5;241m+\u001b[39m\u001b[38;5;241m=\u001b[39m residual  \u001b[38;5;66;03m# Out-of-place addition to prevent in-place errors\u001b[39;00m\n",
      "File \u001b[1;32mc:\\Users\\Hephzibah\\AppData\\Local\\Programs\\Python\\Python310\\lib\\site-packages\\torch\\nn\\modules\\module.py:1736\u001b[0m, in \u001b[0;36mModule._wrapped_call_impl\u001b[1;34m(self, *args, **kwargs)\u001b[0m\n\u001b[0;32m   1734\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_compiled_call_impl(\u001b[38;5;241m*\u001b[39margs, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs)  \u001b[38;5;66;03m# type: ignore[misc]\u001b[39;00m\n\u001b[0;32m   1735\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[1;32m-> 1736\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_call_impl(\u001b[38;5;241m*\u001b[39margs, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs)\n",
      "File \u001b[1;32mc:\\Users\\Hephzibah\\AppData\\Local\\Programs\\Python\\Python310\\lib\\site-packages\\torch\\nn\\modules\\module.py:1747\u001b[0m, in \u001b[0;36mModule._call_impl\u001b[1;34m(self, *args, **kwargs)\u001b[0m\n\u001b[0;32m   1742\u001b[0m \u001b[38;5;66;03m# If we don't have any hooks, we want to skip the rest of the logic in\u001b[39;00m\n\u001b[0;32m   1743\u001b[0m \u001b[38;5;66;03m# this function, and just call forward.\u001b[39;00m\n\u001b[0;32m   1744\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m (\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_backward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_backward_pre_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_pre_hooks\n\u001b[0;32m   1745\u001b[0m         \u001b[38;5;129;01mor\u001b[39;00m _global_backward_pre_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_backward_hooks\n\u001b[0;32m   1746\u001b[0m         \u001b[38;5;129;01mor\u001b[39;00m _global_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_forward_pre_hooks):\n\u001b[1;32m-> 1747\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m forward_call(\u001b[38;5;241m*\u001b[39margs, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs)\n\u001b[0;32m   1749\u001b[0m result \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m\n\u001b[0;32m   1750\u001b[0m called_always_called_hooks \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mset\u001b[39m()\n",
      "File \u001b[1;32mc:\\Users\\Hephzibah\\AppData\\Local\\Programs\\Python\\Python310\\lib\\site-packages\\torch\\nn\\modules\\batchnorm.py:193\u001b[0m, in \u001b[0;36m_BatchNorm.forward\u001b[1;34m(self, input)\u001b[0m\n\u001b[0;32m    186\u001b[0m     bn_training \u001b[38;5;241m=\u001b[39m (\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mrunning_mean \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m) \u001b[38;5;129;01mand\u001b[39;00m (\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mrunning_var \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m)\n\u001b[0;32m    188\u001b[0m \u001b[38;5;250m\u001b[39m\u001b[38;5;124mr\u001b[39m\u001b[38;5;124;03m\"\"\"\u001b[39;00m\n\u001b[0;32m    189\u001b[0m \u001b[38;5;124;03mBuffers are only updated if they are to be tracked and we are in training mode. Thus they only need to be\u001b[39;00m\n\u001b[0;32m    190\u001b[0m \u001b[38;5;124;03mpassed when the update should occur (i.e. in training mode when they are tracked), or when buffer stats are\u001b[39;00m\n\u001b[0;32m    191\u001b[0m \u001b[38;5;124;03mused for normalization (i.e. in eval mode when buffers are not None).\u001b[39;00m\n\u001b[0;32m    192\u001b[0m \u001b[38;5;124;03m\"\"\"\u001b[39;00m\n\u001b[1;32m--> 193\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mF\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mbatch_norm\u001b[49m\u001b[43m(\u001b[49m\n\u001b[0;32m    194\u001b[0m \u001b[43m    \u001b[49m\u001b[38;5;28;43minput\u001b[39;49m\u001b[43m,\u001b[49m\n\u001b[0;32m    195\u001b[0m \u001b[43m    \u001b[49m\u001b[38;5;66;43;03m# If buffers are not to be tracked, ensure that they won't be updated\u001b[39;49;00m\n\u001b[0;32m    196\u001b[0m \u001b[43m    \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mrunning_mean\u001b[49m\n\u001b[0;32m    197\u001b[0m \u001b[43m    \u001b[49m\u001b[38;5;28;43;01mif\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[38;5;129;43;01mnot\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mtraining\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;129;43;01mor\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mtrack_running_stats\u001b[49m\n\u001b[0;32m    198\u001b[0m \u001b[43m    \u001b[49m\u001b[38;5;28;43;01melse\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[38;5;28;43;01mNone\u001b[39;49;00m\u001b[43m,\u001b[49m\n\u001b[0;32m    199\u001b[0m \u001b[43m    \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mrunning_var\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;28;43;01mif\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[38;5;129;43;01mnot\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mtraining\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;129;43;01mor\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mtrack_running_stats\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;28;43;01melse\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[38;5;28;43;01mNone\u001b[39;49;00m\u001b[43m,\u001b[49m\n\u001b[0;32m    200\u001b[0m \u001b[43m    \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mweight\u001b[49m\u001b[43m,\u001b[49m\n\u001b[0;32m    201\u001b[0m \u001b[43m    \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mbias\u001b[49m\u001b[43m,\u001b[49m\n\u001b[0;32m    202\u001b[0m \u001b[43m    \u001b[49m\u001b[43mbn_training\u001b[49m\u001b[43m,\u001b[49m\n\u001b[0;32m    203\u001b[0m \u001b[43m    \u001b[49m\u001b[43mexponential_average_factor\u001b[49m\u001b[43m,\u001b[49m\n\u001b[0;32m    204\u001b[0m \u001b[43m    \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43meps\u001b[49m\u001b[43m,\u001b[49m\n\u001b[0;32m    205\u001b[0m \u001b[43m\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[1;32mc:\\Users\\Hephzibah\\AppData\\Local\\Programs\\Python\\Python310\\lib\\site-packages\\torch\\nn\\functional.py:2812\u001b[0m, in \u001b[0;36mbatch_norm\u001b[1;34m(input, running_mean, running_var, weight, bias, training, momentum, eps)\u001b[0m\n\u001b[0;32m   2809\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m training:\n\u001b[0;32m   2810\u001b[0m     _verify_batch_size(\u001b[38;5;28minput\u001b[39m\u001b[38;5;241m.\u001b[39msize())\n\u001b[1;32m-> 2812\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mtorch\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mbatch_norm\u001b[49m\u001b[43m(\u001b[49m\n\u001b[0;32m   2813\u001b[0m \u001b[43m    \u001b[49m\u001b[38;5;28;43minput\u001b[39;49m\u001b[43m,\u001b[49m\n\u001b[0;32m   2814\u001b[0m \u001b[43m    \u001b[49m\u001b[43mweight\u001b[49m\u001b[43m,\u001b[49m\n\u001b[0;32m   2815\u001b[0m \u001b[43m    \u001b[49m\u001b[43mbias\u001b[49m\u001b[43m,\u001b[49m\n\u001b[0;32m   2816\u001b[0m \u001b[43m    \u001b[49m\u001b[43mrunning_mean\u001b[49m\u001b[43m,\u001b[49m\n\u001b[0;32m   2817\u001b[0m \u001b[43m    \u001b[49m\u001b[43mrunning_var\u001b[49m\u001b[43m,\u001b[49m\n\u001b[0;32m   2818\u001b[0m \u001b[43m    \u001b[49m\u001b[43mtraining\u001b[49m\u001b[43m,\u001b[49m\n\u001b[0;32m   2819\u001b[0m \u001b[43m    \u001b[49m\u001b[43mmomentum\u001b[49m\u001b[43m,\u001b[49m\n\u001b[0;32m   2820\u001b[0m \u001b[43m    \u001b[49m\u001b[43meps\u001b[49m\u001b[43m,\u001b[49m\n\u001b[0;32m   2821\u001b[0m \u001b[43m    \u001b[49m\u001b[43mtorch\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mbackends\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mcudnn\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43menabled\u001b[49m\u001b[43m,\u001b[49m\n\u001b[0;32m   2822\u001b[0m \u001b[43m\u001b[49m\u001b[43m)\u001b[49m\n",
      "\u001b[1;31mKeyboardInterrupt\u001b[0m: "
     ]
    }
   ],
   "source": [
    "# Cell 8: Training Loop with Enhancements\n",
    "\n",
    "# Initialize or load the model\n",
    "net = ConnectNet(num_residual_blocks=hp.num_residual_blocks).to(device)  # Increased residual blocks for deeper network\n",
    "optimizer = optim.Adam(net.parameters(), lr=hp.learning_rate, weight_decay=hp.weight_decay)  # L2 regularization\n",
    "scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='min', factor=0.5, patience=5)\n",
    "\n",
    "start_iteration = 0  # To keep track of iterations when loading\n",
    "\n",
    "# Check if a saved model exists\n",
    "model_filename = 'connect4_model.pth'\n",
    "best_model_filename = 'connect4_best_model.pth'\n",
    "load_existing = False\n",
    "best_eval_score = -math.inf\n",
    "\n",
    "if os.path.exists(model_filename):\n",
    "    user_input = input(\"A saved model was found. Do you want to load it and continue training? (y/n): \")\n",
    "    if user_input.lower() == 'y':\n",
    "        load_existing = True\n",
    "else:\n",
    "    logger.info(\"No saved model found. Training a new model.\")\n",
    "\n",
    "if load_existing:\n",
    "    checkpoint = load_checkpoint(model_filename)\n",
    "    if checkpoint:\n",
    "        net.load_state_dict(checkpoint['model_state_dict'])\n",
    "        optimizer.load_state_dict(checkpoint['optimizer_state_dict'])\n",
    "        scheduler.load_state_dict(checkpoint['scheduler_state_dict'])\n",
    "        start_iteration = checkpoint['iteration']\n",
    "        logger.info(f\"Loaded model from iteration {start_iteration}\")\n",
    "        # Optionally load best_eval_score from checkpoint if stored\n",
    "        if 'best_eval_score' in checkpoint:\n",
    "            best_eval_score = checkpoint['best_eval_score']\n",
    "    else:\n",
    "        logger.warning(\"Failed to load checkpoint. Starting training from scratch.\")\n",
    "else:\n",
    "    logger.info(\"Starting training from scratch.\")\n",
    "\n",
    "total_iterations = hp.total_iterations  # Total number of training iterations you want\n",
    "\n",
    "# Initialize Replay Buffer\n",
    "replay_buffer = ReplayBuffer(max_size=hp.replay_buffer_max_size)  # Increased buffer size\n",
    "\n",
    "for iteration in range(start_iteration, total_iterations):\n",
    "    logger.info(f\"\\n=== Iteration {iteration + 1}/{total_iterations} ===\")\n",
    "    # Self-Play\n",
    "    memory, self_play_stats = self_play(net, num_games=hp.num_self_play_games, num_simulations=hp.num_mcts_simulations)\n",
    "    \n",
    "    # Apply Data Augmentation\n",
    "    augmented_memory = augment_memory(memory)\n",
    "    replay_buffer.add(augmented_memory)\n",
    "    logger.info(f\"Self-Play completed. Stats: {self_play_stats}\")\n",
    "    logger.info(f\"Added {len(augmented_memory)} augmented game states to replay buffer.\")\n",
    "    \n",
    "    # Train the network\n",
    "    avg_value_loss, avg_policy_loss, avg_loss = train(net, replay_buffer, optimizer, \n",
    "                                                      batch_size=hp.batch_size, \n",
    "                                                      epochs=hp.epochs, \n",
    "                                                      clip_grad=hp.clip_grad)\n",
    "    logger.info(f\"Training completed. Avg Value Loss: {avg_value_loss:.4f}, Avg Policy Loss: {avg_policy_loss:.4f}, Avg Total Loss: {avg_loss:.4f}\")\n",
    "    \n",
    "    # Log training metrics\n",
    "    log_training_metrics(iteration + 1, avg_value_loss, avg_policy_loss, avg_loss, self_play_stats)\n",
    "    \n",
    "    # Update the learning rate scheduler\n",
    "    scheduler.step(avg_loss)\n",
    "    \n",
    "    # Periodic Evaluation\n",
    "    eval_stats = None\n",
    "    if (iteration + 1) % hp.evaluation_interval == 0:\n",
    "        eval_stats = evaluate(net, num_games=hp.num_evaluation_games, num_simulations=hp.num_mcts_simulations)\n",
    "        logger.info(f\"Evaluation after {iteration + 1} iterations: {eval_stats}\")\n",
    "        # Log evaluation metrics\n",
    "        log_training_metrics(iteration + 1, avg_value_loss, avg_policy_loss, avg_loss, self_play_stats, eval_stats)\n",
    "        \n",
    "        # Calculate a simple evaluation score (e.g., wins minus losses)\n",
    "        eval_score = eval_stats.get(1, 0) - eval_stats.get(-1, 0)\n",
    "        \n",
    "        # Save the best model\n",
    "        if eval_score > best_eval_score:\n",
    "            best_eval_score = eval_score\n",
    "            save_checkpoint({\n",
    "                'iteration': iteration + 1,\n",
    "                'model_state_dict': net.state_dict(),\n",
    "                'optimizer_state_dict': optimizer.state_dict(),\n",
    "                'scheduler_state_dict': scheduler.state_dict(),\n",
    "                'best_eval_score': best_eval_score\n",
    "            }, filename=best_model_filename)\n",
    "            logger.info(f\"New best model saved at iteration {iteration + 1} with evaluation score {eval_score}\")\n",
    "    \n",
    "    # Save the latest model checkpoint\n",
    "    save_checkpoint({\n",
    "        'iteration': iteration + 1,\n",
    "        'model_state_dict': net.state_dict(),\n",
    "        'optimizer_state_dict': optimizer.state_dict(),\n",
    "        'scheduler_state_dict': scheduler.state_dict(),\n",
    "        'best_eval_score': best_eval_score\n",
    "    }, filename=model_filename)\n",
    "    logger.info(f\"Model checkpoint saved at iteration {iteration + 1}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Playing Against the AI"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "C:\\Users\\Hephzibah\\AppData\\Local\\Temp\\ipykernel_15336\\3456340580.py:47: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.\n",
      "  checkpoint = torch.load(filename, map_location=device)\n",
      "2024-11-22 13:04:24,472 - INFO - Checkpoint loaded from connect4_best_model.pth\n",
      "2024-11-22 13:04:24,506 - INFO - Loaded best model from iteration 10 with evaluation score 6\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Current Board:\n",
      "[[0 0 0 0 0 0 0]\n",
      " [0 0 0 0 0 0 0]\n",
      " [0 0 0 0 0 0 0]\n",
      " [0 0 0 0 0 0 0]\n",
      " [0 0 0 0 0 0 0]\n",
      " [0 0 0 0 0 0 0]]\n",
      "AI selects column 3\n",
      "\n",
      "Current Board:\n",
      "[[ 0  0  0  0  0  0  0]\n",
      " [ 0  0  0  0  0  0  0]\n",
      " [ 0  0  0  0  0  0  0]\n",
      " [ 0  0  0  0  0  0  0]\n",
      " [ 0  0  0 -1  0  0  0]\n",
      " [ 0  0  0  1  0  0  0]]\n",
      "AI selects column 4\n",
      "\n",
      "Current Board:\n",
      "[[ 0  0  0  0  0  0  0]\n",
      " [ 0  0  0  0  0  0  0]\n",
      " [ 0  0  0  0  0  0  0]\n",
      " [ 0  0  0  0  0  0  0]\n",
      " [ 0  0  0 -1  0  0  0]\n",
      " [ 0  0  1  1 -1  0  0]]\n",
      "AI selects column 1\n",
      "\n",
      "Current Board:\n",
      "[[ 0  0  0  0  0  0  0]\n",
      " [ 0  0  0  0  0  0  0]\n",
      " [ 0  0  0  0  0  0  0]\n",
      " [ 0  0  0  0  0  0  0]\n",
      " [ 0  0  0 -1  0  0  0]\n",
      " [ 1 -1  1  1 -1  0  0]]\n",
      "AI selects column 3\n",
      "\n",
      "Current Board:\n",
      "[[ 0  0  0  0  0  0  0]\n",
      " [ 0  0  0  0  0  0  0]\n",
      " [ 0  0  0  0  0  0  0]\n",
      " [ 0  0  0 -1  0  0  0]\n",
      " [ 0  1  0 -1  0  0  0]\n",
      " [ 1 -1  1  1 -1  0  0]]\n",
      "AI selects column 3\n",
      "\n",
      "Current Board:\n",
      "[[ 0  0  0  0  0  0  0]\n",
      " [ 0  0  0 -1  0  0  0]\n",
      " [ 0  0  0  1  0  0  0]\n",
      " [ 0  0  0 -1  0  0  0]\n",
      " [ 0  1  0 -1  0  0  0]\n",
      " [ 1 -1  1  1 -1  0  0]]\n",
      "AI selects column 4\n",
      "\n",
      "Current Board:\n",
      "[[ 0  0  0  0  0  0  0]\n",
      " [ 0  0  0 -1  0  0  0]\n",
      " [ 0  0  0  1  0  0  0]\n",
      " [ 0  0  0 -1  0  0  0]\n",
      " [ 0  1  0 -1 -1  0  0]\n",
      " [ 1 -1  1  1 -1  1  0]]\n",
      "AI selects column 6\n",
      "\n",
      "Current Board:\n",
      "[[ 0  0  0  0  0  0  0]\n",
      " [ 0  0  0 -1  0  0  0]\n",
      " [ 0  0  0  1  0  0  0]\n",
      " [ 0  0  0 -1  1  0  0]\n",
      " [ 0  1  0 -1 -1  0  0]\n",
      " [ 1 -1  1  1 -1  1 -1]]\n",
      "AI selects column 4\n",
      "\n",
      "Current Board:\n",
      "[[ 0  0  0  1  0  0  0]\n",
      " [ 0  0  0 -1  0  0  0]\n",
      " [ 0  0  0  1 -1  0  0]\n",
      " [ 0  0  0 -1  1  0  0]\n",
      " [ 0  1  0 -1 -1  0  0]\n",
      " [ 1 -1  1  1 -1  1 -1]]\n",
      "AI selects column 2\n",
      "\n",
      "Current Board:\n",
      "[[ 0  0  0  1  0  0  0]\n",
      " [ 0  0  0 -1  0  0  0]\n",
      " [ 0  0  0  1 -1  0  0]\n",
      " [ 0  0 -1 -1  1  0  0]\n",
      " [ 0  1  1 -1 -1  0  0]\n",
      " [ 1 -1  1  1 -1  1 -1]]\n",
      "AI selects column 2\n",
      "\n",
      "Current Board:\n",
      "[[ 0  0  0  1  0  0  0]\n",
      " [ 0  0 -1 -1  0  0  0]\n",
      " [ 0  0  1  1 -1  0  0]\n",
      " [ 0  0 -1 -1  1  0  0]\n",
      " [ 0  1  1 -1 -1  0  0]\n",
      " [ 1 -1  1  1 -1  1 -1]]\n",
      "AI selects column 1\n",
      "\n",
      "Current Board:\n",
      "[[ 0  0  0  1  0  0  0]\n",
      " [ 0  0 -1 -1  0  0  0]\n",
      " [ 0  0  1  1 -1  0  0]\n",
      " [ 0 -1 -1 -1  1  0  0]\n",
      " [ 1  1  1 -1 -1  0  0]\n",
      " [ 1 -1  1  1 -1  1 -1]]\n",
      "AI selects column 1\n",
      "\n",
      "Final Board:\n",
      "[[ 0  0  0  1  0  0  0]\n",
      " [ 0  0 -1 -1  0  0  0]\n",
      " [ 0 -1  1  1 -1  0  0]\n",
      " [ 1 -1 -1 -1  1  0  0]\n",
      " [ 1  1  1 -1 -1  0  0]\n",
      " [ 1 -1  1  1 -1  1 -1]]\n",
      "AI wins!\n"
     ]
    }
   ],
   "source": [
    "# Cell 9: Playing Against the AI\n",
    "\n",
    "# Function to Play Against the AI\n",
    "def play_against_ai(net, num_simulations=800):\n",
    "    net.eval()\n",
    "    game = ConnectFour()\n",
    "    while True:\n",
    "        print(\"\\nCurrent Board:\")\n",
    "        print(game.board)\n",
    "        # Human move\n",
    "        try:\n",
    "            move = int(input(\"Enter your move (0-6): \"))\n",
    "        except ValueError:\n",
    "            print(\"Invalid input. Please enter an integer between 0 and 6.\")\n",
    "            continue\n",
    "        if move not in game.valid_moves():\n",
    "            print(\"Invalid move. Try again.\")\n",
    "            continue\n",
    "        game.make_move(move)\n",
    "        winner = game.check_winner()\n",
    "        if winner != 0 or game.is_full():\n",
    "            break\n",
    "        # AI move\n",
    "        root = MCTSNode(game)\n",
    "        ai_move = mcts_search(root, net, num_simulations=num_simulations)\n",
    "        print(f\"AI selects column {ai_move}\")\n",
    "        game.make_move(ai_move)\n",
    "        winner = game.check_winner()\n",
    "        if winner != 0 or game.is_full():\n",
    "            break\n",
    "    print(\"\\nFinal Board:\")\n",
    "    print(game.board)\n",
    "    if winner == 1:\n",
    "        print(\"You win!\")\n",
    "    elif winner == -1:\n",
    "        print(\"AI wins!\")\n",
    "    else:\n",
    "        print(\"It's a draw.\")\n",
    "\n",
    "# Function to Load the Best Trained Model\n",
    "def load_best_model(model_filename='connect4_best_model.pth'):\n",
    "    checkpoint = load_checkpoint(model_filename)\n",
    "    if checkpoint:\n",
    "        net_loaded = ConnectNet(num_residual_blocks=hp.num_residual_blocks).to(device)\n",
    "        net_loaded.load_state_dict(checkpoint['model_state_dict'])\n",
    "        net_loaded.eval()\n",
    "        logger.info(f\"Loaded best model from iteration {checkpoint['iteration']} with evaluation score {checkpoint.get('best_eval_score', 'N/A')}\")\n",
    "        return net_loaded\n",
    "    else:\n",
    "        logger.warning(\"Best model not found.\")\n",
    "        return None\n",
    "    \n",
    "# Load the best trained model\n",
    "best_net = load_best_model()\n",
    "\n",
    "# Play against the AI if the best model was successfully loaded\n",
    "if best_net is not None:\n",
    "    play_against_ai(best_net, num_simulations=hp.num_mcts_simulations)\n",
    "else:\n",
    "    logger.warning(\"Cannot play against AI without a trained best model.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Interactive Widget for training and Playing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "8424e22533a04089ab5b165b08bc11ee",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Button(description='Start Training', style=ButtonStyle())"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "fb83540b38f44b70be728f007cd23d67",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Button(description='Play Against AI', style=ButtonStyle())"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "051405bca9464a2e90517b80a9928782",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Button(button_style='danger', description='Stop Training', style=ButtonStyle())"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "7840e4a2d9a34e0c9a495acffe200347",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Output()"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "8a6632d97f0644af83860646f9e737ba",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Output()"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2024-11-20 22:12:45,188 - INFO - Starting Training...\n",
      "2024-11-20 22:12:45,188 - INFO - \n",
      "=== Iteration 8/1000 ===\n"
     ]
    }
   ],
   "source": [
    "# Cell 10: Interactive Widgets for Training and Playing\n",
    "\n",
    "# Define Buttons\n",
    "train_button = widgets.Button(description=\"Start Training\")\n",
    "play_button = widgets.Button(description=\"Play Against AI\")\n",
    "stop_training_button = widgets.Button(description=\"Stop Training\", button_style='danger')\n",
    "\n",
    "# Define Output Areas\n",
    "train_output = widgets.Output()\n",
    "play_output = widgets.Output()\n",
    "\n",
    "display(train_button, play_button, stop_training_button, train_output, play_output)\n",
    "\n",
    "# Global flag to control training\n",
    "stop_training_flag = False\n",
    "\n",
    "# Define Training Function Wrapper with Threading\n",
    "def start_training_thread():\n",
    "    global stop_training_flag\n",
    "    stop_training_flag = False\n",
    "    with train_output:\n",
    "        clear_output()\n",
    "        logger.info(\"Starting Training...\")\n",
    "    for iteration in range(start_iteration, hp.total_iterations):\n",
    "        if stop_training_flag:\n",
    "            logger.info(\"Training stopped by user.\")\n",
    "            break\n",
    "        logger.info(f\"\\n=== Iteration {iteration + 1}/{hp.total_iterations} ===\")\n",
    "        # Self-Play\n",
    "        memory, self_play_stats = self_play(net, num_games=hp.num_self_play_games, num_simulations=hp.num_mcts_simulations)\n",
    "        \n",
    "        # Apply Data Augmentation\n",
    "        augmented_memory = augment_memory(memory)\n",
    "        replay_buffer.add(augmented_memory)\n",
    "        logger.info(f\"Self-Play completed. Stats: {self_play_stats}\")\n",
    "        logger.info(f\"Added {len(augmented_memory)} augmented game states to replay buffer.\")\n",
    "        \n",
    "        # Train the network\n",
    "        avg_value_loss, avg_policy_loss, avg_loss = train(net, replay_buffer, optimizer, \n",
    "                                                          batch_size=hp.batch_size, \n",
    "                                                          epochs=hp.epochs, \n",
    "                                                          clip_grad=hp.clip_grad)\n",
    "        logger.info(f\"Training completed. Avg Value Loss: {avg_value_loss:.4f}, Avg Policy Loss: {avg_policy_loss:.4f}, Avg Total Loss: {avg_loss:.4f}\")\n",
    "        \n",
    "        # Log training metrics\n",
    "        log_training_metrics(iteration + 1, avg_value_loss, avg_policy_loss, avg_loss, self_play_stats)\n",
    "        \n",
    "        # Update the learning rate scheduler\n",
    "        scheduler.step(avg_loss)\n",
    "        \n",
    "        # Periodic Evaluation\n",
    "        eval_stats = None\n",
    "        if (iteration + 1) % hp.evaluation_interval == 0:\n",
    "            eval_stats = evaluate(net, num_games=hp.num_evaluation_games, num_simulations=hp.num_mcts_simulations)\n",
    "            logger.info(f\"Evaluation after {iteration + 1} iterations: {eval_stats}\")\n",
    "            # Log evaluation metrics\n",
    "            log_training_metrics(iteration + 1, avg_value_loss, avg_policy_loss, avg_loss, self_play_stats, eval_stats)\n",
    "            \n",
    "            # Calculate a simple evaluation score (e.g., wins minus losses)\n",
    "            eval_score = eval_stats.get(1, 0) - eval_stats.get(-1, 0)\n",
    "            \n",
    "            # Save the best model\n",
    "            if eval_score > best_eval_score:\n",
    "                best_eval_score = eval_score\n",
    "                save_checkpoint({\n",
    "                    'iteration': iteration + 1,\n",
    "                    'model_state_dict': net.state_dict(),\n",
    "                    'optimizer_state_dict': optimizer.state_dict(),\n",
    "                    'scheduler_state_dict': scheduler.state_dict(),\n",
    "                    'best_eval_score': best_eval_score\n",
    "                }, filename=best_model_filename)\n",
    "                logger.info(f\"New best model saved at iteration {iteration + 1} with evaluation score {eval_score}\")\n",
    "        \n",
    "        # Save the latest model checkpoint\n",
    "        save_checkpoint({\n",
    "            'iteration': iteration + 1,\n",
    "            'model_state_dict': net.state_dict(),\n",
    "            'optimizer_state_dict': optimizer.state_dict(),\n",
    "            'scheduler_state_dict': scheduler.state_dict(),\n",
    "            'best_eval_score': best_eval_score\n",
    "        }, filename=model_filename)\n",
    "        logger.info(f\"Model checkpoint saved at iteration {iteration + 1}\")\n",
    "    \n",
    "    logger.info(\"Training Completed.\")\n",
    "\n",
    "# Define Play Function Wrapper\n",
    "def start_playing(b):\n",
    "    with play_output:\n",
    "        clear_output()\n",
    "        logger.info(\"Loading Best Model...\")\n",
    "        loaded_net = load_best_model()\n",
    "        if loaded_net is not None:\n",
    "            play_against_ai(loaded_net, num_simulations=hp.num_mcts_simulations)\n",
    "        else:\n",
    "            logger.warning(\"Cannot play against AI without a trained best model.\")\n",
    "\n",
    "# Define Stop Training Function\n",
    "def stop_training_func(b):\n",
    "    global stop_training_flag\n",
    "    stop_training_flag = True\n",
    "    with train_output:\n",
    "        logger.info(\"Stopping Training...\")\n",
    "\n",
    "# Assign Functions to Buttons\n",
    "def on_train_button_clicked(b):\n",
    "    training_thread = threading.Thread(target=start_training_thread)\n",
    "    training_thread.start()\n",
    "\n",
    "train_button.on_click(on_train_button_clicked)\n",
    "play_button.on_click(start_playing)\n",
    "stop_training_button.on_click(stop_training_func)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "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.10.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}