8.ConnectFour.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "2d090f5e",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1.23.2\n",
      "1.12.1+cu102\n"
     ]
    }
   ],
   "source": [
    "import numpy as np\n",
    "print(np.__version__)\n",
    "\n",
    "\n",
    "import torch\n",
    "print(torch.__version__)\n",
    "\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "\n",
    "torch.manual_seed(0)\n",
    "\n",
    "from tqdm.notebook import trange\n",
    "\n",
    "import random\n",
    "import math"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "a097e1b5",
   "metadata": {},
   "outputs": [],
   "source": [
    "class TicTacToe:\n",
    "    def __init__(self):\n",
    "        self.row_count = 3\n",
    "        self.column_count = 3\n",
    "        self.action_size = self.row_count * self.column_count\n",
    "        \n",
    "    def __repr__(self):\n",
    "        return \"TicTacToe\"\n",
    "        \n",
    "    def get_initial_state(self):\n",
    "        return np.zeros((self.row_count, self.column_count))\n",
    "    \n",
    "    def get_next_state(self, state, action, player):\n",
    "        row = action // self.column_count\n",
    "        column = action % self.column_count\n",
    "        state[row, column] = player\n",
    "        return state\n",
    "    \n",
    "    def get_valid_moves(self, state):\n",
    "        return (state.reshape(-1) == 0).astype(np.uint8)\n",
    "    \n",
    "    def check_win(self, state, action):\n",
    "        if action == None:\n",
    "            return False\n",
    "        \n",
    "        row = action // self.column_count\n",
    "        column = action % self.column_count\n",
    "        player = state[row, column]\n",
    "        \n",
    "        return (\n",
    "            np.sum(state[row, :]) == player * self.column_count\n",
    "            or np.sum(state[:, column]) == player * self.row_count\n",
    "            or np.sum(np.diag(state)) == player * self.row_count\n",
    "            or np.sum(np.diag(np.flip(state, axis=0))) == player * self.row_count\n",
    "        )\n",
    "    \n",
    "    def get_value_and_terminated(self, state, action):\n",
    "        if self.check_win(state, action):\n",
    "            return 1, True\n",
    "        if np.sum(self.get_valid_moves(state)) == 0:\n",
    "            return 0, True\n",
    "        return 0, False\n",
    "    \n",
    "    def get_opponent(self, player):\n",
    "        return -player\n",
    "    \n",
    "    def get_opponent_value(self, value):\n",
    "        return -value\n",
    "    \n",
    "    def change_perspective(self, state, player):\n",
    "        return state * player\n",
    "    \n",
    "    def get_encoded_state(self, state):\n",
    "        encoded_state = np.stack(\n",
    "            (state == -1, state == 0, state == 1)\n",
    "        ).astype(np.float32)\n",
    "        \n",
    "        return encoded_state"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "682c4ca4",
   "metadata": {},
   "outputs": [],
   "source": [
    "class ConnectFour:\n",
    "    def __init__(self):\n",
    "        self.row_count = 6\n",
    "        self.column_count = 7\n",
    "        self.action_size = self.column_count\n",
    "        self.in_a_row = 4\n",
    "        \n",
    "    def __repr__(self):\n",
    "        return \"ConnectFour\"\n",
    "        \n",
    "    def get_initial_state(self):\n",
    "        return np.zeros((self.row_count, self.column_count))\n",
    "    \n",
    "    def get_next_state(self, state, action, player):\n",
    "        row = np.max(np.where(state[:, action] == 0))\n",
    "        state[row, action] = player\n",
    "        return state\n",
    "    \n",
    "    def get_valid_moves(self, state):\n",
    "        return (state[0] == 0).astype(np.uint8)\n",
    "    \n",
    "    def check_win(self, state, action):\n",
    "        if action == None:\n",
    "            return False\n",
    "        \n",
    "        row = np.min(np.where(state[:, action] != 0))\n",
    "        column = action\n",
    "        player = state[row][column]\n",
    "\n",
    "        def count(offset_row, offset_column):\n",
    "            for i in range(1, self.in_a_row):\n",
    "                r = row + offset_row * i\n",
    "                c = action + offset_column * i\n",
    "                if (\n",
    "                    r < 0 \n",
    "                    or r >= self.row_count\n",
    "                    or c < 0 \n",
    "                    or c >= self.column_count\n",
    "                    or state[r][c] != player\n",
    "                ):\n",
    "                    return i - 1\n",
    "            return self.in_a_row - 1\n",
    "\n",
    "        return (\n",
    "            count(1, 0) >= self.in_a_row - 1 # vertical\n",
    "            or (count(0, 1) + count(0, -1)) >= self.in_a_row - 1 # horizontal\n",
    "            or (count(1, 1) + count(-1, -1)) >= self.in_a_row - 1 # top left diagonal\n",
    "            or (count(1, -1) + count(-1, 1)) >= self.in_a_row - 1 # top right diagonal\n",
    "        )\n",
    "    \n",
    "    def get_value_and_terminated(self, state, action):\n",
    "        if self.check_win(state, action):\n",
    "            return 1, True\n",
    "        if np.sum(self.get_valid_moves(state)) == 0:\n",
    "            return 0, True\n",
    "        return 0, False\n",
    "    \n",
    "    def get_opponent(self, player):\n",
    "        return -player\n",
    "    \n",
    "    def get_opponent_value(self, value):\n",
    "        return -value\n",
    "    \n",
    "    def change_perspective(self, state, player):\n",
    "        return state * player\n",
    "    \n",
    "    def get_encoded_state(self, state):\n",
    "        encoded_state = np.stack(\n",
    "            (state == -1, state == 0, state == 1)\n",
    "        ).astype(np.float32)\n",
    "        \n",
    "        return encoded_state"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "02e5b58b",
   "metadata": {},
   "outputs": [],
   "source": [
    "class ResNet(nn.Module):\n",
    "    def __init__(self, game, num_resBlocks, num_hidden, device):\n",
    "        super().__init__()\n",
    "        \n",
    "        self.device = device\n",
    "        self.startBlock = nn.Sequential(\n",
    "            nn.Conv2d(3, num_hidden, kernel_size=3, padding=1),\n",
    "            nn.BatchNorm2d(num_hidden),\n",
    "            nn.ReLU()\n",
    "        )\n",
    "        \n",
    "        self.backBone = nn.ModuleList(\n",
    "            [ResBlock(num_hidden) for i in range(num_resBlocks)]\n",
    "        )\n",
    "        \n",
    "        self.policyHead = nn.Sequential(\n",
    "            nn.Conv2d(num_hidden, 32, kernel_size=3, padding=1),\n",
    "            nn.BatchNorm2d(32),\n",
    "            nn.ReLU(),\n",
    "            nn.Flatten(),\n",
    "            nn.Linear(32 * game.row_count * game.column_count, game.action_size)\n",
    "        )\n",
    "        \n",
    "        self.valueHead = nn.Sequential(\n",
    "            nn.Conv2d(num_hidden, 3, kernel_size=3, padding=1),\n",
    "            nn.BatchNorm2d(3),\n",
    "            nn.ReLU(),\n",
    "            nn.Flatten(),\n",
    "            nn.Linear(3 * game.row_count * game.column_count, 1),\n",
    "            nn.Tanh()\n",
    "        )\n",
    "        \n",
    "        self.to(device)\n",
    "        \n",
    "    def forward(self, x):\n",
    "        x = self.startBlock(x)\n",
    "        for resBlock in self.backBone:\n",
    "            x = resBlock(x)\n",
    "        policy = self.policyHead(x)\n",
    "        value = self.valueHead(x)\n",
    "        return policy, value\n",
    "        \n",
    "        \n",
    "class ResBlock(nn.Module):\n",
    "    def __init__(self, num_hidden):\n",
    "        super().__init__()\n",
    "        self.conv1 = nn.Conv2d(num_hidden, num_hidden, kernel_size=3, padding=1)\n",
    "        self.bn1 = nn.BatchNorm2d(num_hidden)\n",
    "        self.conv2 = nn.Conv2d(num_hidden, num_hidden, kernel_size=3, padding=1)\n",
    "        self.bn2 = nn.BatchNorm2d(num_hidden)\n",
    "        \n",
    "    def forward(self, x):\n",
    "        residual = x\n",
    "        x = F.relu(self.bn1(self.conv1(x)))\n",
    "        x = self.bn2(self.conv2(x))\n",
    "        x += residual\n",
    "        x = F.relu(x)\n",
    "        return x\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "797c3205",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0.10280606150627136\n",
      "[[ 0.  0. -1.]\n",
      " [ 0. -1.  0.]\n",
      " [ 1.  0.  1.]]\n",
      "tensor([[[[0., 0., 1.],\n",
      "          [0., 1., 0.],\n",
      "          [0., 0., 0.]],\n",
      "\n",
      "         [[1., 1., 0.],\n",
      "          [1., 0., 1.],\n",
      "          [0., 1., 0.]],\n",
      "\n",
      "         [[0., 0., 0.],\n",
      "          [0., 0., 0.],\n",
      "          [1., 0., 1.]]]], device='cuda:0')\n"
     ]
    },
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAiwAAAGdCAYAAAAxCSikAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjUuMSwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy/YYfK9AAAACXBIWXMAAA9hAAAPYQGoP6dpAAAhtElEQVR4nO3df0xd9f3H8RdQ+dFacJYUbEu9VTspFqGFwqCmuIxIt26Kc4iNswybGpNSqXdhltpCTNVbXUvoCitjWV2MEljjWqvr2PAqVVcqFto5/NGa/RDS5l4gbqB0QsO93z/89rprb39cxN4Pl+cjudGe+zmH98nV9Jlzz+WGuN1utwAAAAwWGugBAAAALoZgAQAAxiNYAACA8QgWAABgPIIFAAAYj2ABAADGI1gAAIDxCBYAAGC8KYEeYDy4XC6dOnVK06dPV0hISKDHAQAAl8DtduuTTz7RrFmzFBp64WsoQREsp06dUkJCQqDHAAAAY9DT06M5c+ZccE1QBMv06dMlfX7C0dHRAZ4GAABcisHBQSUkJHj+Hr+QoAiWs28DRUdHEywAAEwwl3I7BzfdAgAA4xEsAADAeAQLAAAwHsECAACMR7AAAADjESwAAMB4BAsAADAewQIAAIxHsAAAAOMRLAAAwHgECwAAMN6YgqW2tlYWi0WRkZHKzMxUe3v7ede+++67uuuuu2SxWBQSEqLq6upz1thsNi1ZskTTp0/XzJkzlZ+fr+PHj49lNAAAEIT8DpampiZZrVZVVlaqs7NTKSkpysvLU29vr8/1p0+f1nXXXaetW7cqPj7e55qDBw9q7dq1Onz4sFpaWnTmzBnddtttGhoa8nc8AAAQhELcbrfbnx0yMzO1ZMkS1dTUSJJcLpcSEhK0bt06bdiw4YL7WiwWrV+/XuvXr7/gur6+Ps2cOVMHDx7UsmXLLjrT4OCgYmJiNDAwwLc1AwAwQfjz9/cUfw48MjKijo4OlZeXe7aFhoYqNzdXbW1tY5vWh4GBAUnS1Vdf7fP54eFhDQ8Pe/48ODg4bj8bADAxWDb8IdAjXNS/tq4I9AhBw6+3hPr7+zU6Oqq4uDiv7XFxcXI4HOMykMvl0vr167V06VItXLjQ5xqbzaaYmBjPIyEhYVx+NgAAMJNxnxJau3aturq61NjYeN415eXlGhgY8Dx6enou44QAAOBy8+stodjYWIWFhcnpdHptdzqd572h1h8lJSV6+eWX9frrr2vOnDnnXRcREaGIiIiv/PMAAMDE4FewhIeHKy0tTXa7Xfn5+ZI+fwvHbrerpKRkzEO43W6tW7dOe/fuVWtrq+bNmzfmY+HCeM8XADAR+RUskmS1WlVUVKT09HRlZGSourpaQ0NDKi4uliStWrVKs2fPls1mk/T5jbrvvfee599PnjypY8eO6corr9QNN9wg6fO3gRoaGvTiiy9q+vTpnvthYmJiFBUVNS4nCgAAJi6/g6WwsFB9fX2qqKiQw+FQamqqmpubPTfidnd3KzT0i1tjTp06pUWLFnn+vG3bNm3btk05OTlqbW2VJO3atUuSdOutt3r9rGeeeUY/+clP/B0RkwRXiwBg8vA7WKTP7zU531tAZyPkLIvFoov9qhc/fxUMAACYZIz7lBAAAMCXESwAAMB4BAsAADAewQIAAIxHsAAAAOMRLAAAwHgECwAAMB7BAgAAjEewAAAA4xEsAADAeAQLAAAwHsECAACMR7AAAADjjenbmicby4Y/BHqEi/rX1hWBHgEAgK8NV1gAAIDxuMICAECAcSX/4rjCAgAAjEewAAAA4xEsAADAeAQLAAAwHsECAACMR7AAAADjESwAAMB4BAsAADAewQIAAIxHsAAAAOMRLAAAwHgECwAAMB7BAgAAjMe3NQPABfAtuoAZuMICAACMR7AAAADj8ZYQAEwSvL2FiYwrLAAAwHgECwAAMB7BAgAAjEewAAAA4xEsAADAeAQLAAAwHsECAACMx+9hAQzB78gAgPPjCgsAADAewQIAAIxHsAAAAOMRLAAAwHgECwAAMB7BAgAAjDemYKmtrZXFYlFkZKQyMzPV3t5+3rXvvvuu7rrrLlksFoWEhKi6uvorHxMAAEwufgdLU1OTrFarKisr1dnZqZSUFOXl5am3t9fn+tOnT+u6667T1q1bFR8fPy7HBAAAk4vfwVJVVaU1a9aouLhYSUlJqqur09SpU7V7926f65csWaKf//znuueeexQRETEuxwQAAJOLX8EyMjKijo4O5ebmfnGA0FDl5uaqra1tTAOM5ZjDw8MaHBz0egAAgODlV7D09/drdHRUcXFxXtvj4uLkcDjGNMBYjmmz2RQTE+N5JCQkjOlnAwCAiWFCfkqovLxcAwMDnkdPT0+gRwIAAF8jv778MDY2VmFhYXI6nV7bnU7neW+o/TqOGRERcd77YQAEHl/kCGC8+XWFJTw8XGlpabLb7Z5tLpdLdrtdWVlZYxrg6zgmAAAILn5dYZEkq9WqoqIipaenKyMjQ9XV1RoaGlJxcbEkadWqVZo9e7ZsNpukz2+qfe+99zz/fvLkSR07dkxXXnmlbrjhhks6JgAAmNz8DpbCwkL19fWpoqJCDodDqampam5u9tw0293drdDQLy7cnDp1SosWLfL8edu2bdq2bZtycnLU2tp6SccEAACTm9/BIkklJSUqKSnx+dzZCDnLYrHI7XZ/pWMCAIDJbUJ+SggAAEwuBAsAADAewQIAAIxHsAAAAOMRLAAAwHgECwAAMB7BAgAAjEewAAAA4xEsAADAeAQLAAAwHsECAACMR7AAAADjESwAAMB4BAsAADAewQIAAIxHsAAAAOMRLAAAwHgECwAAMB7BAgAAjEewAAAA4xEsAADAeAQLAAAwHsECAACMR7AAAADjESwAAMB4BAsAADAewQIAAIxHsAAAAOMRLAAAwHgECwAAMB7BAgAAjEewAAAA4xEsAADAeAQLAAAwHsECAACMR7AAAADjESwAAMB4BAsAADAewQIAAIxHsAAAAOMRLAAAwHgECwAAMB7BAgAAjEewAAAA4xEsAADAeAQLAAAwHsECAACMR7AAAADjjSlYamtrZbFYFBkZqczMTLW3t19w/Z49e5SYmKjIyEglJyfrwIEDXs9/+umnKikp0Zw5cxQVFaWkpCTV1dWNZTQAABCE/A6WpqYmWa1WVVZWqrOzUykpKcrLy1Nvb6/P9YcOHdLKlSu1evVqHT16VPn5+crPz1dXV5dnjdVqVXNzs5577jm9//77Wr9+vUpKSrR///6xnxkAAAgafgdLVVWV1qxZo+LiYs+VkKlTp2r37t0+1+/YsUPLly9XWVmZFixYoC1btmjx4sWqqanxrDl06JCKiop06623ymKx6IEHHlBKSspFr9wAAIDJwa9gGRkZUUdHh3Jzc784QGiocnNz1dbW5nOftrY2r/WSlJeX57U+Oztb+/fv18mTJ+V2u/Xaa6/pxIkTuu2223wec3h4WIODg14PAAAQvPwKlv7+fo2OjiouLs5re1xcnBwOh899HA7HRdfv3LlTSUlJmjNnjsLDw7V8+XLV1tZq2bJlPo9ps9kUExPjeSQkJPhzGgAAYIIx4lNCO3fu1OHDh7V//351dHRo+/btWrt2rV555RWf68vLyzUwMOB59PT0XOaJAQDA5TTFn8WxsbEKCwuT0+n02u50OhUfH+9zn/j4+Auu/+9//6uNGzdq7969WrFihSTp5ptv1rFjx7Rt27Zz3k6SpIiICEVERPgzOgAAmMD8usISHh6utLQ02e12zzaXyyW73a6srCyf+2RlZXmtl6SWlhbP+jNnzujMmTMKDfUeJSwsTC6Xy5/xAABAkPLrCov0+UeQi4qKlJ6eroyMDFVXV2toaEjFxcWSpFWrVmn27Nmy2WySpNLSUuXk5Gj79u1asWKFGhsbdeTIEdXX10uSoqOjlZOTo7KyMkVFRenaa6/VwYMH9eyzz6qqqmocTxUAAExUfgdLYWGh+vr6VFFRIYfDodTUVDU3N3turO3u7va6WpKdna2GhgZt2rRJGzdu1Pz587Vv3z4tXLjQs6axsVHl5eW699579fHHH+vaa6/VE088oQcffHAcThEAAEx0fgeLJJWUlKikpMTnc62tredsKygoUEFBwXmPFx8fr2eeeWYsowAAgEnAiE8JAQAAXAjBAgAAjEewAAAA4xEsAADAeAQLAAAwHsECAACMR7AAAADjESwAAMB4BAsAADAewQIAAIxHsAAAAOMRLAAAwHgECwAAMB7BAgAAjEewAAAA4xEsAADAeAQLAAAwHsECAACMR7AAAADjESwAAMB4BAsAADAewQIAAIxHsAAAAOMRLAAAwHgECwAAMB7BAgAAjEewAAAA4xEsAADAeAQLAAAwHsECAACMR7AAAADjESwAAMB4BAsAADAewQIAAIxHsAAAAOMRLAAAwHgECwAAMB7BAgAAjEewAAAA4xEsAADAeAQLAAAwHsECAACMR7AAAADjESwAAMB4BAsAADAewQIAAIxHsAAAAOONKVhqa2tlsVgUGRmpzMxMtbe3X3D9nj17lJiYqMjISCUnJ+vAgQPnrHn//fd1++23KyYmRtOmTdOSJUvU3d09lvEAAECQ8TtYmpqaZLVaVVlZqc7OTqWkpCgvL0+9vb0+1x86dEgrV67U6tWrdfToUeXn5ys/P19dXV2eNX//+991yy23KDExUa2trXrnnXe0efNmRUZGjv3MAABA0PA7WKqqqrRmzRoVFxcrKSlJdXV1mjp1qnbv3u1z/Y4dO7R8+XKVlZVpwYIF2rJlixYvXqyamhrPmkcffVTf+9739PTTT2vRokW6/vrrdfvtt2vmzJljPzMAABA0/AqWkZERdXR0KDc394sDhIYqNzdXbW1tPvdpa2vzWi9JeXl5nvUul0t/+MMf9M1vflN5eXmaOXOmMjMztW/fPj9PBQAABCu/gqW/v1+jo6OKi4vz2h4XFyeHw+FzH4fDccH1vb29+vTTT7V161YtX75cf/7zn3XnnXfqhz/8oQ4ePOjzmMPDwxocHPR6AACA4DUl0AO4XC5J0h133KGHH35YkpSamqpDhw6prq5OOTk55+xjs9n02GOPXdY5AQBA4Ph1hSU2NlZhYWFyOp1e251Op+Lj433uEx8ff8H1sbGxmjJlipKSkrzWLFiw4LyfEiovL9fAwIDn0dPT489pAACACcavYAkPD1daWprsdrtnm8vlkt1uV1ZWls99srKyvNZLUktLi2d9eHi4lixZouPHj3utOXHihK699lqfx4yIiFB0dLTXAwAABC+/3xKyWq0qKipSenq6MjIyVF1draGhIRUXF0uSVq1apdmzZ8tms0mSSktLlZOTo+3bt2vFihVqbGzUkSNHVF9f7zlmWVmZCgsLtWzZMn37299Wc3OzXnrpJbW2to7PWQIAgAnN72ApLCxUX1+fKioq5HA4lJqaqubmZs+Ntd3d3QoN/eLCTXZ2thoaGrRp0yZt3LhR8+fP1759+7Rw4ULPmjvvvFN1dXWy2Wx66KGHdOONN+qFF17QLbfcMg6nCAAAJrox3XRbUlKikpISn8/5uipSUFCggoKCCx7z/vvv1/333z+WcQAAQJDju4QAAIDxCBYAAGA8ggUAABiPYAEAAMYjWAAAgPEIFgAAYDyCBQAAGI9gAQAAxiNYAACA8QgWAABgPIIFAAAYj2ABAADGI1gAAIDxCBYAAGA8ggUAABiPYAEAAMYjWAAAgPEIFgAAYDyCBQAAGI9gAQAAxiNYAACA8QgWAABgPIIFAAAYj2ABAADGI1gAAIDxCBYAAGA8ggUAABiPYAEAAMYjWAAAgPEIFgAAYDyCBQAAGI9gAQAAxiNYAACA8QgWAABgPIIFAAAYj2ABAADGI1gAAIDxCBYAAGA8ggUAABiPYAEAAMYjWAAAgPEIFgAAYDyCBQAAGI9gAQAAxiNYAACA8QgWAABgPIIFAAAYj2ABAADGG1Ow1NbWymKxKDIyUpmZmWpvb7/g+j179igxMVGRkZFKTk7WgQMHzrv2wQcfVEhIiKqrq8cyGgAACEJ+B0tTU5OsVqsqKyvV2dmplJQU5eXlqbe31+f6Q4cOaeXKlVq9erWOHj2q/Px85efnq6ur65y1e/fu1eHDhzVr1iz/zwQAAAQtv4OlqqpKa9asUXFxsZKSklRXV6epU6dq9+7dPtfv2LFDy5cvV1lZmRYsWKAtW7Zo8eLFqqmp8Vp38uRJrVu3Ts8//7yuuOKKsZ0NAAAISn4Fy8jIiDo6OpSbm/vFAUJDlZubq7a2Np/7tLW1ea2XpLy8PK/1LpdL9913n8rKynTTTTdddI7h4WENDg56PQAAQPDyK1j6+/s1OjqquLg4r+1xcXFyOBw+93E4HBdd/9RTT2nKlCl66KGHLmkOm82mmJgYzyMhIcGf0wAAABNMwD8l1NHRoR07dui3v/2tQkJCLmmf8vJyDQwMeB49PT1f85QAACCQ/AqW2NhYhYWFyel0em13Op2Kj4/3uU98fPwF17/xxhvq7e3V3LlzNWXKFE2ZMkUfffSRfvrTn8pisfg8ZkREhKKjo70eAAAgePkVLOHh4UpLS5Pdbvdsc7lcstvtysrK8rlPVlaW13pJamlp8ay/77779M477+jYsWOex6xZs1RWVqY//elP/p4PAAAIQlP83cFqtaqoqEjp6enKyMhQdXW1hoaGVFxcLElatWqVZs+eLZvNJkkqLS1VTk6Otm/frhUrVqixsVFHjhxRfX29JGnGjBmaMWOG18+44oorFB8frxtvvPGrnh8AAAgCfgdLYWGh+vr6VFFRIYfDodTUVDU3N3turO3u7lZo6BcXbrKzs9XQ0KBNmzZp48aNmj9/vvbt26eFCxeO31kAAICg5newSFJJSYlKSkp8Ptfa2nrOtoKCAhUUFFzy8f/1r3+NZSwAABCkAv4pIQAAgIshWAAAgPEIFgAAYDyCBQAAGI9gAQAAxiNYAACA8QgWAABgPIIFAAAYj2ABAADGI1gAAIDxCBYAAGA8ggUAABiPYAEAAMYjWAAAgPEIFgAAYDyCBQAAGI9gAQAAxiNYAACA8QgWAABgPIIFAAAYj2ABAADGI1gAAIDxCBYAAGA8ggUAABiPYAEAAMYjWAAAgPEIFgAAYDyCBQAAGI9gAQAAxiNYAACA8QgWAABgPIIFAAAYj2ABAADGI1gAAIDxCBYAAGA8ggUAABiPYAEAAMYjWAAAgPEIFgAAYDyCBQAAGI9gAQAAxiNYAACA8QgWAABgPIIFAAAYj2ABAADGI1gAAIDxCBYAAGA8ggUAABhvTMFSW1sri8WiyMhIZWZmqr29/YLr9+zZo8TEREVGRio5OVkHDhzwPHfmzBk98sgjSk5O1rRp0zRr1iytWrVKp06dGstoAAAgCPkdLE1NTbJaraqsrFRnZ6dSUlKUl5en3t5en+sPHTqklStXavXq1Tp69Kjy8/OVn5+vrq4uSdLp06fV2dmpzZs3q7OzU7///e91/Phx3X777V/tzAAAQNDwO1iqqqq0Zs0aFRcXKykpSXV1dZo6dap2797tc/2OHTu0fPlylZWVacGCBdqyZYsWL16smpoaSVJMTIxaWlp0991368Ybb9S3vvUt1dTUqKOjQ93d3V/t7AAAQFDwK1hGRkbU0dGh3NzcLw4QGqrc3Fy1tbX53Ketrc1rvSTl5eWdd70kDQwMKCQkRFdddZXP54eHhzU4OOj1AAAAwcuvYOnv79fo6Kji4uK8tsfFxcnhcPjcx+Fw+LX+s88+0yOPPKKVK1cqOjra5xqbzaaYmBjPIyEhwZ/TAAAAE4xRnxI6c+aM7r77brndbu3ateu868rLyzUwMOB59PT0XMYpAQDA5TbFn8WxsbEKCwuT0+n02u50OhUfH+9zn/j4+EtafzZWPvroI7366qvnvboiSREREYqIiPBndAAAMIH5dYUlPDxcaWlpstvtnm0ul0t2u11ZWVk+98nKyvJaL0ktLS1e68/GyocffqhXXnlFM2bM8GcsAAAQ5Py6wiJJVqtVRUVFSk9PV0ZGhqqrqzU0NKTi4mJJ0qpVqzR79mzZbDZJUmlpqXJycrR9+3atWLFCjY2NOnLkiOrr6yV9His/+tGP1NnZqZdfflmjo6Oe+1uuvvpqhYeHj9e5AgCACcrvYCksLFRfX58qKirkcDiUmpqq5uZmz4213d3dCg394sJNdna2GhoatGnTJm3cuFHz58/Xvn37tHDhQknSyZMntX//fklSamqq18967bXXdOutt47x1AAAQLDwO1gkqaSkRCUlJT6fa21tPWdbQUGBCgoKfK63WCxyu91jGQMAAEwSRn1KCAAAwBeCBQAAGI9gAQAAxiNYAACA8QgWAABgPIIFAAAYj2ABAADGI1gAAIDxCBYAAGA8ggUAABiPYAEAAMYjWAAAgPEIFgAAYDyCBQAAGI9gAQAAxiNYAACA8QgWAABgPIIFAAAYj2ABAADGI1gAAIDxCBYAAGA8ggUAABiPYAEAAMYjWAAAgPEIFgAAYDyCBQAAGI9gAQAAxiNYAACA8QgWAABgPIIFAAAYj2ABAADGI1gAAIDxCBYAAGA8ggUAABiPYAEAAMYjWAAAgPEIFgAAYDyCBQAAGI9gAQAAxiNYAACA8QgWAABgPIIFAAAYj2ABAADGI1gAAIDxCBYAAGA8ggUAABiPYAEAAMYbU7DU1tbKYrEoMjJSmZmZam9vv+D6PXv2KDExUZGRkUpOTtaBAwe8nne73aqoqNA111yjqKgo5ebm6sMPPxzLaAAAIAj5HSxNTU2yWq2qrKxUZ2enUlJSlJeXp97eXp/rDx06pJUrV2r16tU6evSo8vPzlZ+fr66uLs+ap59+Wr/4xS9UV1ent956S9OmTVNeXp4+++yzsZ8ZAAAIGn4HS1VVldasWaPi4mIlJSWprq5OU6dO1e7du32u37Fjh5YvX66ysjItWLBAW7Zs0eLFi1VTUyPp86sr1dXV2rRpk+644w7dfPPNevbZZ3Xq1Cnt27fvK50cAAAIDlP8WTwyMqKOjg6Vl5d7toWGhio3N1dtbW0+92lra5PVavXalpeX54mRf/7zn3I4HMrNzfU8HxMTo8zMTLW1temee+4555jDw8MaHh72/HlgYECSNDg46M/pXDLX8Omv5bjj6VLPnXO5vPz5bzKYzodzubwm47lIwXU+wXQuYzmm2+2+6Fq/gqW/v1+jo6OKi4vz2h4XF6cPPvjA5z4Oh8PneofD4Xn+7Lbzrfkym82mxx577JztCQkJl3YiQSimOtATjB/OxVzBdD6ci5mC6Vyk4Dqfr/NcPvnkE8XExFxwjV/BYory8nKvqzYul0sff/yxZsyYoZCQkABOdmkGBweVkJCgnp4eRUdHB3oc/D9eFzPxupiL18ZME+l1cbvd+uSTTzRr1qyLrvUrWGJjYxUWFian0+m13el0Kj4+3uc+8fHxF1x/9p9Op1PXXHON15rU1FSfx4yIiFBERITXtquuusqfUzFCdHS08f8xTUa8LmbidTEXr42ZJsrrcrErK2f5ddNteHi40tLSZLfbPdtcLpfsdruysrJ87pOVleW1XpJaWlo86+fNm6f4+HivNYODg3rrrbfOe0wAADC5+P2WkNVqVVFRkdLT05WRkaHq6moNDQ2puLhYkrRq1SrNnj1bNptNklRaWqqcnBxt375dK1asUGNjo44cOaL6+npJUkhIiNavX6/HH39c8+fP17x587R582bNmjVL+fn543emAABgwvI7WAoLC9XX16eKigo5HA6lpqaqubnZc9Nsd3e3QkO/uHCTnZ2thoYGbdq0SRs3btT8+fO1b98+LVy40LPmZz/7mYaGhvTAAw/oP//5j2655RY1NzcrMjJyHE7RPBEREaqsrDznbS0EFq+LmXhdzMVrY6ZgfV1C3JfyWSIAAIAA4ruEAACA8QgWAABgPIIFAAAYj2ABAADGI1gus9raWlksFkVGRiozM1Pt7e2BHmnSs9lsWrJkiaZPn66ZM2cqPz9fx48fD/RY+JKtW7d6fg0CAu/kyZP68Y9/rBkzZigqKkrJyck6cuRIoMea1EZHR7V582bNmzdPUVFRuv7667Vly5ZL+p6eiYBguYyamppktVpVWVmpzs5OpaSkKC8vT729vYEebVI7ePCg1q5dq8OHD6ulpUVnzpzRbbfdpqGhoUCPhv/39ttv61e/+pVuvvnmQI8CSf/+97+1dOlSXXHFFfrjH/+o9957T9u3b9c3vvGNQI82qT311FPatWuXampq9P777+upp57S008/rZ07dwZ6tHHBx5ovo8zMTC1ZskQ1NTWSPv8twQkJCVq3bp02bNgQ4OlwVl9fn2bOnKmDBw9q2bJlgR5n0vv000+1ePFi/fKXv9Tjjz+u1NRUVVdXB3qsSW3Dhg36y1/+ojfeeCPQo+B/fP/731dcXJx+85vfeLbdddddioqK0nPPPRfAycYHV1guk5GREXV0dCg3N9ezLTQ0VLm5uWprawvgZPiygYEBSdLVV18d4EkgSWvXrtWKFSu8/t9BYO3fv1/p6ekqKCjQzJkztWjRIv36178O9FiTXnZ2tux2u06cOCFJ+utf/6o333xT3/3udwM82fiYkN/WPBH19/drdHTU8xuBz4qLi9MHH3wQoKnwZS6XS+vXr9fSpUu9fhszAqOxsVGdnZ16++23Az0K/sc//vEP7dq1S1arVRs3btTbb7+thx56SOHh4SoqKgr0eJPWhg0bNDg4qMTERIWFhWl0dFRPPPGE7r333kCPNi4IFuB/rF27Vl1dXXrzzTcDPcqk19PTo9LSUrW0tATt13RMVC6XS+np6XryySclSYsWLVJXV5fq6uoIlgD63e9+p+eff14NDQ266aabdOzYMa1fv16zZs0KiteFYLlMYmNjFRYWJqfT6bXd6XQqPj4+QFPhf5WUlOjll1/W66+/rjlz5gR6nEmvo6NDvb29Wrx4sWfb6OioXn/9ddXU1Gh4eFhhYWEBnHDyuuaaa5SUlOS1bcGCBXrhhRcCNBEkqaysTBs2bNA999wjSUpOTtZHH30km80WFMHCPSyXSXh4uNLS0mS32z3bXC6X7Ha7srKyAjgZ3G63SkpKtHfvXr366quaN29eoEeCpO985zv629/+pmPHjnke6enpuvfee3Xs2DFiJYCWLl16zkf/T5w4oWuvvTZAE0GSTp8+7fXlw5IUFhYml8sVoInGF1dYLiOr1aqioiKlp6crIyND1dXVGhoaUnFxcaBHm9TWrl2rhoYGvfjii5o+fbocDockKSYmRlFRUQGebvKaPn36OfcRTZs2TTNmzOD+ogB7+OGHlZ2drSeffFJ333232tvbVV9fr/r6+kCPNqn94Ac/0BNPPKG5c+fqpptu0tGjR1VVVaX7778/0KONDzcuq507d7rnzp3rDg8Pd2dkZLgPHz4c6JEmPUk+H88880ygR8OX5OTkuEtLSwM9Btxu90svveReuHChOyIiwp2YmOiur68P9EiT3uDgoLu0tNQ9d+5cd2RkpPu6665zP/roo+7h4eFAjzYu+D0sAADAeNzDAgAAjEewAAAA4xEsAADAeAQLAAAwHsECAACMR7AAAADjESwAAMB4BAsAADAewQIAAIxHsAAAAOMRLAAAwHgECwAAMN7/AQzOwbPRd90MAAAAAElFTkSuQmCC",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "\n",
    "tictactoe = TicTacToe()\n",
    "\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "\n",
    "state = tictactoe.get_initial_state()\n",
    "state = tictactoe.get_next_state(state, 2, -1)\n",
    "state = tictactoe.get_next_state(state, 4, -1)\n",
    "state = tictactoe.get_next_state(state, 6, 1)\n",
    "state = tictactoe.get_next_state(state, 8, 1)\n",
    "\n",
    "\n",
    "encoded_state = tictactoe.get_encoded_state(state)\n",
    "\n",
    "tensor_state = torch.tensor(encoded_state, device=device).unsqueeze(0)\n",
    "\n",
    "model = ResNet(tictactoe, 4, 64, device=device)\n",
    "# model.load_state_dict(torch.load('model_2.pt', map_location=device))\n",
    "model.eval()\n",
    "\n",
    "policy, value = model(tensor_state)\n",
    "value = value.item()\n",
    "policy = torch.softmax(policy, axis=1).squeeze(0).detach().cpu().numpy()\n",
    "\n",
    "print(value)\n",
    "\n",
    "print(state)\n",
    "print(tensor_state)\n",
    "\n",
    "plt.bar(range(tictactoe.action_size), policy)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "21866526",
   "metadata": {},
   "outputs": [],
   "source": [
    "class Node:\n",
    "    def __init__(self, game, args, state, parent=None, action_taken=None, prior=0, visit_count=0):\n",
    "        self.game = game\n",
    "        self.args = args\n",
    "        self.state = state\n",
    "        self.parent = parent\n",
    "        self.action_taken = action_taken\n",
    "        self.prior = prior\n",
    "        \n",
    "        self.children = []\n",
    "        \n",
    "        self.visit_count = visit_count\n",
    "        self.value_sum = 0\n",
    "        \n",
    "    def is_fully_expanded(self):\n",
    "        return len(self.children) > 0\n",
    "    \n",
    "    def select(self):\n",
    "        best_child = None\n",
    "        best_ucb = -np.inf\n",
    "        \n",
    "        for child in self.children:\n",
    "            ucb = self.get_ucb(child)\n",
    "            if ucb > best_ucb:\n",
    "                best_child = child\n",
    "                best_ucb = ucb\n",
    "                \n",
    "        return best_child\n",
    "    \n",
    "    def get_ucb(self, child):\n",
    "        if child.visit_count == 0:\n",
    "            q_value = 0\n",
    "        else:\n",
    "            q_value = 1 - ((child.value_sum / child.visit_count) + 1) / 2\n",
    "        return q_value + self.args['C'] * (math.sqrt(self.visit_count) / (child.visit_count + 1)) * child.prior\n",
    "    \n",
    "    def expand(self, policy):\n",
    "        for action, prob in enumerate(policy):\n",
    "            if prob > 0:\n",
    "                child_state = self.state.copy()\n",
    "                child_state = self.game.get_next_state(child_state, action, 1)\n",
    "                child_state = self.game.change_perspective(child_state, player=-1)\n",
    "\n",
    "                child = Node(self.game, self.args, child_state, self, action, prob)\n",
    "                self.children.append(child)\n",
    "                \n",
    "        return child\n",
    "            \n",
    "    def backpropagate(self, value):\n",
    "        self.value_sum += value\n",
    "        self.visit_count += 1\n",
    "        \n",
    "        value = self.game.get_opponent_value(value)\n",
    "        if self.parent is not None:\n",
    "            self.parent.backpropagate(value)  \n",
    "\n",
    "\n",
    "class MCTS:\n",
    "    def __init__(self, game, args, model):\n",
    "        self.game = game\n",
    "        self.args = args\n",
    "        self.model = model\n",
    "        \n",
    "    @torch.no_grad()\n",
    "    def search(self, state):\n",
    "        root = Node(self.game, self.args, state, visit_count=1)\n",
    "        \n",
    "        policy, _ = self.model(\n",
    "            torch.tensor(self.game.get_encoded_state(state), device=self.model.device).unsqueeze(0)\n",
    "        )\n",
    "        policy = torch.softmax(policy, axis=1).squeeze(0).cpu().numpy()\n",
    "        policy = (1 - self.args['dirichlet_epsilon']) * policy + self.args['dirichlet_epsilon'] \\\n",
    "            * np.random.dirichlet([self.args['dirichlet_alpha']] * self.game.action_size)\n",
    "        \n",
    "        valid_moves = self.game.get_valid_moves(state)\n",
    "        policy *= valid_moves\n",
    "        policy /= np.sum(policy)\n",
    "        root.expand(policy)\n",
    "        \n",
    "        for search in range(self.args['num_searches']):\n",
    "            node = root\n",
    "            \n",
    "            while node.is_fully_expanded():\n",
    "                node = node.select()\n",
    "                \n",
    "            value, is_terminal = self.game.get_value_and_terminated(node.state, node.action_taken)\n",
    "            value = self.game.get_opponent_value(value)\n",
    "            \n",
    "            if not is_terminal:\n",
    "                policy, value = self.model(\n",
    "                    torch.tensor(self.game.get_encoded_state(node.state), device=self.model.device).unsqueeze(0)\n",
    "                )\n",
    "                policy = torch.softmax(policy, axis=1).squeeze(0).cpu().numpy()\n",
    "                valid_moves = self.game.get_valid_moves(node.state)\n",
    "                policy *= valid_moves\n",
    "                policy /= np.sum(policy)\n",
    "                \n",
    "                value = value.item()\n",
    "                \n",
    "                node.expand(policy)\n",
    "                \n",
    "            node.backpropagate(value)    \n",
    "            \n",
    "            \n",
    "        action_probs = np.zeros(self.game.action_size)\n",
    "        for child in root.children:\n",
    "            action_probs[child.action_taken] = child.visit_count\n",
    "        action_probs /= np.sum(action_probs)\n",
    "        return action_probs\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "a3b28ab8",
   "metadata": {},
   "outputs": [],
   "source": [
    "class AlphaZero:\n",
    "    def __init__(self, model, optimizer, game, args):\n",
    "        self.model = model\n",
    "        self.optimizer = optimizer\n",
    "        self.game = game\n",
    "        self.args = args\n",
    "        self.mcts = MCTS(game, args, model)\n",
    "        \n",
    "    def selfPlay(self):\n",
    "        memory = []\n",
    "        player = 1\n",
    "        state = self.game.get_initial_state()\n",
    "        \n",
    "        while True:\n",
    "            neutral_state = self.game.change_perspective(state, player)\n",
    "            action_probs = self.mcts.search(neutral_state)\n",
    "            \n",
    "            memory.append((neutral_state, action_probs, player))\n",
    "            \n",
    "            temperature_action_probs = action_probs ** (1 / self.args['temperature'])\n",
    "            action = np.random.choice(self.game.action_size, p=action_probs) # change to temperature_action_probs\n",
    "            \n",
    "            state = self.game.get_next_state(state, action, player)\n",
    "            \n",
    "            value, is_terminal = self.game.get_value_and_terminated(state, action)\n",
    "            \n",
    "            if is_terminal:\n",
    "                returnMemory = []\n",
    "                for hist_neutral_state, hist_action_probs, hist_player in memory:\n",
    "                    hist_outcome = value if hist_player == player else self.game.get_opponent_value(value)\n",
    "                    returnMemory.append((\n",
    "                        self.game.get_encoded_state(hist_neutral_state),\n",
    "                        hist_action_probs,\n",
    "                        hist_outcome\n",
    "                    ))\n",
    "                return returnMemory\n",
    "            \n",
    "            player = self.game.get_opponent(player)\n",
    "                \n",
    "    def train(self, memory):\n",
    "        random.shuffle(memory)\n",
    "        for batchIdx in range(0, len(memory), self.args['batch_size']):\n",
    "            sample = memory[batchIdx:min(len(memory) - 1, batchIdx + self.args['batch_size'])] # Change to memory[batchIdx:batchIdx+self.args['batch_size']] in case of an error\n",
    "            state, policy_targets, value_targets = zip(*sample)\n",
    "            \n",
    "            state, policy_targets, value_targets = np.array(state), np.array(policy_targets), np.array(value_targets).reshape(-1, 1)\n",
    "            \n",
    "            state = torch.tensor(state, dtype=torch.float32, device=self.model.device)\n",
    "            policy_targets = torch.tensor(policy_targets, dtype=torch.float32, device=self.model.device)\n",
    "            value_targets = torch.tensor(value_targets, dtype=torch.float32, device=self.model.device)\n",
    "            \n",
    "            out_policy, out_value = self.model(state)\n",
    "            \n",
    "            policy_loss = F.cross_entropy(out_policy, policy_targets)\n",
    "            value_loss = F.mse_loss(out_value, value_targets)\n",
    "            loss = policy_loss + value_loss\n",
    "            \n",
    "            optimizer.zero_grad() # change to self.optimizer\n",
    "            loss.backward()\n",
    "            optimizer.step() # change to self.optimizer\n",
    "    \n",
    "    def learn(self):\n",
    "        for iteration in range(self.args['num_iterations']):\n",
    "            memory = []\n",
    "            \n",
    "            self.model.eval()\n",
    "            for selfPlay_iteration in trange(self.args['num_selfPlay_iterations']):\n",
    "                memory += self.selfPlay()\n",
    "                \n",
    "            self.model.train()\n",
    "            for epoch in trange(self.args['num_epochs']):\n",
    "                self.train(memory)\n",
    "            \n",
    "            torch.save(self.model.state_dict(), f\"model_{iteration}_{self.game}.pt\")\n",
    "            torch.save(self.optimizer.state_dict(), f\"optimizer_{iteration}_{self.game}.pt\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "24bd91ef",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "f8c0b18d070d4ec8a11bb33eb7ab7168",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/5 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "9ac704e993e949e5ac9d30cb51c584fd",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/4 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "8aa8432f8d5b46f9bb34d2c55660ef18",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/5 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "game = ConnectFour()\n",
    "\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "\n",
    "model = ResNet(game, 9, 128, device)\n",
    "\n",
    "optimizer = torch.optim.Adam(model.parameters(), lr=0.001, weight_decay=0.0001)\n",
    "\n",
    "args = {\n",
    "    'C': 2,\n",
    "    'num_searches': 600,\n",
    "    'num_iterations': 8,\n",
    "    'num_selfPlay_iterations': 500,\n",
    "    'num_parallel_games': 100,\n",
    "    'num_epochs': 4,\n",
    "    'batch_size': 128,\n",
    "    'temperature': 1.25,\n",
    "    'dirichlet_epsilon': 0.25,\n",
    "    'dirichlet_alpha': 0.3\n",
    "}\n",
    "\n",
    "alphaZero = AlphaZero(model, optimizer, game, args)\n",
    "alphaZero.learn()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "7c470145",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[[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",
      "valid_moves [0, 1, 2, 3, 4, 5, 6]\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",
      " [1. 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",
      " [ 1. -1.  0.  0.  0.  0.  0.]]\n",
      "valid_moves [0, 1, 2, 3, 4, 5, 6]\n"
     ]
    },
    {
     "ename": "ValueError",
     "evalue": "invalid literal for int() with base 10: ''",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mValueError\u001b[0m                                Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[11], line 28\u001b[0m\n\u001b[1;32m     26\u001b[0m valid_moves \u001b[39m=\u001b[39m game\u001b[39m.\u001b[39mget_valid_moves(state)\n\u001b[1;32m     27\u001b[0m \u001b[39mprint\u001b[39m(\u001b[39m\"\u001b[39m\u001b[39mvalid_moves\u001b[39m\u001b[39m\"\u001b[39m, [i \u001b[39mfor\u001b[39;00m i \u001b[39min\u001b[39;00m \u001b[39mrange\u001b[39m(game\u001b[39m.\u001b[39maction_size) \u001b[39mif\u001b[39;00m valid_moves[i] \u001b[39m==\u001b[39m \u001b[39m1\u001b[39m])\n\u001b[0;32m---> 28\u001b[0m action \u001b[39m=\u001b[39m \u001b[39mint\u001b[39;49m(\u001b[39minput\u001b[39;49m(\u001b[39mf\u001b[39;49m\u001b[39m\"\u001b[39;49m\u001b[39m{\u001b[39;49;00mplayer\u001b[39m}\u001b[39;49;00m\u001b[39m:\u001b[39;49m\u001b[39m\"\u001b[39;49m))\n\u001b[1;32m     30\u001b[0m \u001b[39mif\u001b[39;00m valid_moves[action] \u001b[39m==\u001b[39m \u001b[39m0\u001b[39m:\n\u001b[1;32m     31\u001b[0m     \u001b[39mprint\u001b[39m(\u001b[39m\"\u001b[39m\u001b[39maction not valid\u001b[39m\u001b[39m\"\u001b[39m)\n",
      "\u001b[0;31mValueError\u001b[0m: invalid literal for int() with base 10: ''"
     ]
    }
   ],
   "source": [
    "game = ConnectFour()\n",
    "player = 1\n",
    "\n",
    "args = {\n",
    "    'C': 2,\n",
    "    'num_searches': 100,\n",
    "    'dirichlet_epsilon': 0.,\n",
    "    'dirichlet_alpha': 0.3\n",
    "}\n",
    "\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "\n",
    "model = ResNet(game, 9, 128, device)\n",
    "# model.load_state_dict(torch.load(\"model_0_ConnectFour.pt\", map_location=device))\n",
    "model.eval()\n",
    "\n",
    "mcts = MCTS(game, args, model)\n",
    "\n",
    "state = game.get_initial_state()\n",
    "\n",
    "\n",
    "while True:\n",
    "    print(state)\n",
    "    \n",
    "    if player == 1:\n",
    "        valid_moves = game.get_valid_moves(state)\n",
    "        print(\"valid_moves\", [i for i in range(game.action_size) if valid_moves[i] == 1])\n",
    "        action = int(input(f\"{player}:\"))\n",
    "\n",
    "        if valid_moves[action] == 0:\n",
    "            print(\"action not valid\")\n",
    "            continue\n",
    "            \n",
    "    else:\n",
    "        neutral_state = game.change_perspective(state, player)\n",
    "        mcts_probs = mcts.search(neutral_state)\n",
    "        action = np.argmax(mcts_probs)\n",
    "        \n",
    "    state = game.get_next_state(state, action, player)\n",
    "    \n",
    "    value, is_terminal = game.get_value_and_terminated(state, action)\n",
    "    \n",
    "    if is_terminal:\n",
    "        print(state)\n",
    "        if value == 1:\n",
    "            print(player, \"won\")\n",
    "        else:\n",
    "            print(\"draw\")\n",
    "        break\n",
    "        \n",
    "    player = game.get_opponent(player)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "myenv",
   "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.8.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}