alphazero_tutorial.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1cPlGyQRjgc-"
      },
      "source": [
        "# AlphaZero Tutorial: TicTacToe / Connect4"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Le4eKmv5rwHp"
      },
      "outputs": [],
      "source": [
        "\"\"\"\n",
        "AlphaZero for TicTacToe and Connect4\n",
        "https://github.com/rlglab/rlg-tutorial\n",
        "\n",
        "Reinforcement Learning and Games (RLG) Lab, IIS, Academia Sinica, Taiwan\n",
        "https://rlg.iis.sinica.edu.tw\n",
        "\n",
        "References:\n",
        "[1] D. Silver et al., \"A general reinforcement learning algorithm that masters chess, shogi, and Go through self-play,\" Science 362, 2018.\n",
        "\"\"\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "IDOYs0ARrwHq"
      },
      "source": [
        "## Required Components"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Xm5sdbc-rwHq"
      },
      "source": [
        "### Dependences"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "QvPYt2Y6v4Dy"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "import copy\n",
        "import shutil\n",
        "import numpy as np\n",
        "import torch\n",
        "import torch.nn as nn\n",
        "from typing import Tuple\n",
        "from datetime import datetime\n",
        "from dataclasses import dataclass"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "W5tFJykc7r0_"
      },
      "source": [
        "### Device Setup"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "hX-SWpcU7sm6"
      },
      "outputs": [],
      "source": [
        "if torch.cuda.is_available():\n",
        "    DEVICE = torch.device('cuda')\n",
        "    !nvidia-smi\n",
        "    print()\n",
        "else:\n",
        "    DEVICE = torch.device('cpu')\n",
        "print('CPU')\n",
        "!cat /proc/cpuinfo | grep 'processor\\|model\\ name'\n",
        "print()\n",
        "!python3 --version\n",
        "print('torch version   :', torch.__version__)\n",
        "print('use device      :', DEVICE)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ccV4smYBxTbj"
      },
      "source": [
        "### Environments"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "kgHek7TPwTnX"
      },
      "outputs": [],
      "source": [
        "class TicTacToe:\n",
        "    BOARD_HEIGHT = 3\n",
        "    BOARD_WIDTH = 3\n",
        "\n",
        "    def __init__(self):\n",
        "        self.board: np.ndarray = np.zeros(0)\n",
        "        self.action_history = []\n",
        "        self.current_player: float = 0.\n",
        "        self.winner: float = 0.\n",
        "\n",
        "    def reset(self) -> None:\n",
        "        self.board = np.zeros(self.BOARD_HEIGHT * self.BOARD_WIDTH, dtype=np.single)\n",
        "        self.action_history.clear()\n",
        "        self.current_player = 1.\n",
        "        self.winner = 0.\n",
        "\n",
        "    def act(self, action: int) -> None:\n",
        "        if self.board[action] == 0.:\n",
        "            self.board[action] = self.current_player\n",
        "            self.action_history.append(action)\n",
        "            self._update_winner()\n",
        "            self.current_player *= -1.\n",
        "        else:\n",
        "            raise ValueError('invalid action id.')\n",
        "\n",
        "    def get_legal_actions(self) -> list:\n",
        "        return list(np.where(self.board == 0.)[0])\n",
        "\n",
        "    def is_terminal(self) -> bool:\n",
        "        return self.winner != 0. or not np.any(self.board == 0.)\n",
        "\n",
        "    def get_eval_score(self) -> float:\n",
        "        return self.winner\n",
        "\n",
        "    def get_features(self) -> np.ndarray:\n",
        "        \"\"\"\n",
        "        4 channels\n",
        "            0. own position\n",
        "            1. opponent position\n",
        "            2. Nought turn\n",
        "            3. Cross turn\n",
        "        \"\"\"\n",
        "        features = []\n",
        "        for channel in range(4):\n",
        "            if channel == 0:\n",
        "                features.append(np.where(self.board == self.current_player, 1., 0.))\n",
        "            elif channel == 1:\n",
        "                features.append(np.where(self.board == self.current_player * -1., 1., 0.))\n",
        "            elif channel == 2:\n",
        "                features.append(np.ones_like(self.board) if self.current_player == 1. else np.zeros_like(self.board))\n",
        "            elif channel == 3:\n",
        "                features.append(np.ones_like(self.board) if self.current_player == -1. else np.zeros_like(self.board))\n",
        "        return np.stack(features, dtype=np.single).reshape((-1, self.BOARD_HEIGHT, self.BOARD_WIDTH))\n",
        "\n",
        "    @staticmethod\n",
        "    def get_num_input_channels() -> int:\n",
        "        return 4\n",
        "\n",
        "    def get_input_channel_height(self) -> int:\n",
        "        return self.BOARD_HEIGHT\n",
        "\n",
        "    def get_input_channel_width(self) -> int:\n",
        "        return self.BOARD_WIDTH\n",
        "\n",
        "    def get_policy_size(self) -> int:\n",
        "        return self.BOARD_HEIGHT * self.BOARD_WIDTH\n",
        "\n",
        "    def to_string(self) -> str:\n",
        "        result = np.empty(self.board.shape, dtype=str)\n",
        "        result[np.where(self.board == 0.)] = ' '\n",
        "        result[np.where(self.board == 1.)] = 'O'\n",
        "        result[np.where(self.board == -1.)] = 'X'\n",
        "        return str(result.reshape((self.BOARD_HEIGHT, self.BOARD_WIDTH)))\n",
        "\n",
        "    @staticmethod\n",
        "    def _get_all_line() -> Tuple:\n",
        "        return [0, 1, 2], [3, 4, 5], [6, 7, 8], [0, 3, 6], [1, 4, 7], [2, 5, 8], [0, 4, 8], [2, 4, 6]\n",
        "\n",
        "    def _update_winner(self) -> None:\n",
        "        for line in self._get_all_line():\n",
        "            line_values = self.board[line]\n",
        "            if np.all(line_values != 0.) and np.all(line_values == self.current_player):\n",
        "                self.winner = self.current_player\n",
        "                self.current_player = 0.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "dpmBmnEZkhxt"
      },
      "outputs": [],
      "source": [
        "class ConnectFour(TicTacToe):\n",
        "    BOARD_HEIGHT = 6\n",
        "    BOARD_WIDTH = 7\n",
        "\n",
        "    def act(self, action: int) -> None:\n",
        "        is_empty = (self.board[action] == 0.)\n",
        "        has_below = (action < self.BOARD_WIDTH or self.board[action - self.BOARD_WIDTH] != 0.)\n",
        "        if is_empty and has_below and not self.is_terminal():\n",
        "            self.board[action] = self.current_player\n",
        "            self.action_history.append(action)\n",
        "            num_connection = self._check_connected(action)\n",
        "            if num_connection >= 4:\n",
        "                self.winner = self.current_player\n",
        "                self.current_player = 0.\n",
        "            else:\n",
        "                self.current_player *= -1.\n",
        "        else:\n",
        "            raise ValueError('invalid action id.')\n",
        "\n",
        "    def get_legal_actions(self) -> list:\n",
        "        candidate_actions = set(np.where(self.board == 0.)[0])\n",
        "        legal_actions = []\n",
        "        for width in range(self.BOARD_WIDTH):\n",
        "            for action in range(width, self.get_policy_size(), self.BOARD_WIDTH):\n",
        "                if action in candidate_actions:\n",
        "                    legal_actions.append(action)\n",
        "                    break\n",
        "        legal_actions.sort()\n",
        "        return legal_actions\n",
        "\n",
        "    def to_string(self) -> str:\n",
        "        result = np.empty(self.board.shape, dtype=str)\n",
        "        result[np.where(self.board == 1.)] = 'O'\n",
        "        result[np.where(self.board == -1.)] = 'X'\n",
        "        result[np.where(self.board == 0.)] = ' '\n",
        "        return str(np.flip(result.reshape(self.BOARD_HEIGHT, self.BOARD_WIDTH), 0))\n",
        "\n",
        "    def _check_connected(self, action: int) -> int:\n",
        "        num_connection = 1 + max(\n",
        "            self._get_num_connection(action, -1, -1) + self._get_num_connection(action, 1, 1),\n",
        "            self._get_num_connection(action, 1, -1) + self._get_num_connection(action, -1, 1),\n",
        "            self._get_num_connection(action, -1, 0) + self._get_num_connection(action, 1, 0),\n",
        "            self._get_num_connection(action, 0, -1) + self._get_num_connection(action, 0, 1)\n",
        "        )\n",
        "        return num_connection\n",
        "\n",
        "    def _get_num_connection(self, action: int, row_shift: int, column_shift: int) -> int:\n",
        "        player = self.board[action]\n",
        "        row = action // self.BOARD_WIDTH\n",
        "        column = action % self.BOARD_WIDTH\n",
        "        num_connection = 0\n",
        "        while True:\n",
        "            row += row_shift\n",
        "            column += column_shift\n",
        "            if not (0 <= row < self.BOARD_HEIGHT and 0 <= column < self.BOARD_WIDTH):\n",
        "                break\n",
        "            shift_action = row * self.BOARD_WIDTH + column\n",
        "            if self.board[shift_action] != player:\n",
        "                break\n",
        "            num_connection += 1\n",
        "        return num_connection"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "bO4HLm_hxzs1"
      },
      "outputs": [],
      "source": [
        "def test_env(env_class) -> None:\n",
        "    seeds = [100, 200, 300]\n",
        "    for seed in seeds:\n",
        "        np.random.seed(seed)\n",
        "        env = env_class()\n",
        "        env.reset()\n",
        "        while not env.is_terminal():\n",
        "            legal_actions = env.get_legal_actions()\n",
        "            action_id = np.random.choice(legal_actions)\n",
        "            print('player:', env.current_player, '(O)' if env.current_player == 1.0 else '(X)')\n",
        "            print('board:')\n",
        "            print(env.to_string())\n",
        "            print('features:')\n",
        "            print(env.get_features())\n",
        "            env.act(action_id)\n",
        "            print('-----------------------')\n",
        "            print('act action id:', action_id)\n",
        "\n",
        "        print('score:', env.get_eval_score())\n",
        "        print(env.to_string())\n",
        "        print('========================')\n",
        "\n",
        "test_env(TicTacToe)\n",
        "test_env(ConnectFour)\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qXQEhaARmKju"
      },
      "source": [
        "### MCTS"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "kTzL4EgfmLO8"
      },
      "outputs": [],
      "source": [
        "class Node:\n",
        "    def __init__(self, player: float, action: int, policy: float = 0):\n",
        "        self.player: float = player\n",
        "        self.action: int = action\n",
        "        self.policy: float = policy\n",
        "        self.visit_count: int = 0\n",
        "        self.mean_value: float = 0.\n",
        "        self.children: list[Node] = []\n",
        "\n",
        "    def is_leaf(self) -> bool:\n",
        "        return len(self.children) == 0\n",
        "\n",
        "    def add_child(self, player: float, action: int, policy: float) -> None:\n",
        "        child_node = Node(player, action, policy)\n",
        "        self.children.append(child_node)\n",
        "\n",
        "    def update(self, value: float) -> None:\n",
        "        self.visit_count += 1\n",
        "        self.mean_value += (value - self.mean_value) / self.visit_count\n",
        "\n",
        "\n",
        "class MCTS:\n",
        "    DIRICHLET_NOISE_ALPHA = 0.3  # Usually (1 / sqrt(number of actions))\n",
        "    DIRICHLET_NOISE_EPSILON = 0.25\n",
        "    PUCT_C1 = 1.25\n",
        "    PUCT_C2 = 19652\n",
        "    TEMPERATURE = 1.0\n",
        "\n",
        "    def __init__(self, env: TicTacToe, network: nn.Module):\n",
        "        self.env: TicTacToe = env\n",
        "        self.network: nn.Module = network\n",
        "        self.root: Node = Node(0, 0)\n",
        "\n",
        "    def simulate(self, num_simulations: int, use_dirichlet_noise: bool = True) -> None:\n",
        "        if self.root.visit_count == 0:\n",
        "            node_path = [self.root]\n",
        "            self.expand_and_evaluate(node_path)\n",
        "            if use_dirichlet_noise:\n",
        "                self.add_exploration_noise(self.root)\n",
        "\n",
        "        for _ in range(num_simulations):\n",
        "            # Selection\n",
        "            node_path = self.selection()\n",
        "            # Expansion and evaluation\n",
        "            value = self.expand_and_evaluate(node_path)\n",
        "            # Backup\n",
        "            self.backup(node_path, value)\n",
        "\n",
        "    def selection(self) -> list[Node]:\n",
        "        node_path = [self.root]\n",
        "        node = self.root\n",
        "        while not node.is_leaf():\n",
        "            node = self.select_child(node)\n",
        "            node_path.append(node)\n",
        "        return node_path\n",
        "\n",
        "    def expand_and_evaluate(self, node_path) -> float:\n",
        "        # go to state\n",
        "        env_transition = copy.deepcopy(self.env)\n",
        "        for child in node_path[1:]:\n",
        "            env_transition.act(child.action)\n",
        "\n",
        "        if env_transition.is_terminal():\n",
        "            return env_transition.get_eval_score()\n",
        "\n",
        "        features = env_transition.get_features()\n",
        "        features_tensor = torch.from_numpy(features).unsqueeze(0).to(DEVICE)\n",
        "        _, policy, value = self.network(features_tensor)\n",
        "        policy = policy.squeeze(0)\n",
        "        player = env_transition.current_player\n",
        "        leaf_node = node_path[-1]\n",
        "        for action in env_transition.get_legal_actions():\n",
        "            leaf_node.add_child(player, action, policy[action].item())\n",
        "        return value.squeeze(0).item()\n",
        "\n",
        "    @staticmethod\n",
        "    def backup(node_path: list[Node], value: float) -> None:\n",
        "        for node in node_path:\n",
        "            node.update(value * node.player)\n",
        "\n",
        "    def decide_action(self, use_softmax: bool = True) -> int:\n",
        "        candidate_actions = []\n",
        "        action_weights = []\n",
        "        for child in self.root.children:\n",
        "            candidate_actions.append(child.action)\n",
        "            action_weights.append(child.visit_count ** (1 / self.TEMPERATURE))\n",
        "        if np.sum(action_weights) == 0:\n",
        "            action_weights = np.ones_like(action_weights, dtype=np.single)\n",
        "        action_weights /= np.sum(action_weights)\n",
        "        if use_softmax:\n",
        "            selected_action = np.random.choice(candidate_actions, p=action_weights)\n",
        "        else:\n",
        "            selected_action = candidate_actions[np.argmax(action_weights)]\n",
        "        return selected_action\n",
        "\n",
        "    def add_exploration_noise(self, parent: Node) -> None:\n",
        "        dirichlet_noise = np.random.gamma(self.DIRICHLET_NOISE_ALPHA, 1, len(parent.children))\n",
        "        for child, noise in zip(parent.children, dirichlet_noise):\n",
        "            child.policy = child.policy * (1 - self.DIRICHLET_NOISE_EPSILON) + noise * self.DIRICHLET_NOISE_EPSILON\n",
        "\n",
        "    def select_child(self, parent: Node) -> Node:\n",
        "        # ============== TODO ==============\n",
        "        # hint: select the child with the highest PUCT score\n",
        "        # hint: self.PUCT_C1 and self.PUCT_C2 are PUCT constants\n",
        "        best_child = np.random.choice(parent.children)\n",
        "        return best_child\n",
        "\n",
        "    def get_normalize_child_visits(self) -> np.ndarray:\n",
        "        child_visits = np.zeros(self.env.get_policy_size(), dtype=np.single)\n",
        "        for child in self.root.children:\n",
        "            child_visits[child.action] = child.visit_count / self.root.visit_count\n",
        "        return child_visits\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qdanapkXnzep"
      },
      "source": [
        "### Replay Buffer"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "OiJtymUHnzw9"
      },
      "outputs": [],
      "source": [
        "@dataclass\n",
        "class Game:\n",
        "    action_history: list[int]\n",
        "    terminal_value: float\n",
        "    child_visits: list[np.ndarray]\n",
        "\n",
        "\n",
        "class ReplayBuffer:\n",
        "    def __init__(self, buffer_num_games: int):\n",
        "        self.buffer: list[Game] = []\n",
        "        self.buffer_num_games: int = buffer_num_games\n",
        "\n",
        "    def save_game(self, game: Game) -> None:\n",
        "        if len(self.buffer) >= self.buffer_num_games:\n",
        "            self.buffer.pop(0)\n",
        "        self.buffer.append(game)\n",
        "\n",
        "    def sample_batch(self, env: TicTacToe, batch_size: int) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:\n",
        "        features = []\n",
        "        target_value = []\n",
        "        target_policy = []\n",
        "        select_weights = np.array([len(game.action_history) for game in self.buffer])\n",
        "        select_weights = select_weights / np.sum(select_weights)\n",
        "        sample_games_index = np.random.choice(len(self.buffer), size=batch_size, p=select_weights, replace=True)\n",
        "        unique, counts = np.unique(sample_games_index, return_counts=True)\n",
        "        for game_index, count in zip(unique, counts):\n",
        "            env.reset()\n",
        "            game = self.buffer[game_index]\n",
        "            env_steps = np.random.randint(len(game.action_history), size=count)\n",
        "            env_steps.sort()\n",
        "            for step in env_steps:\n",
        "                for action in game.action_history[len(env.action_history): step]:\n",
        "                    env.act(action)\n",
        "                features.append(env.get_features())\n",
        "                target_value.append(game.terminal_value)\n",
        "                target_policy.append(game.child_visits[step])\n",
        "\n",
        "        features = np.stack(features, dtype=np.single)\n",
        "        target_policy = np.stack(target_policy, dtype=np.single)\n",
        "        target_value = np.stack(target_value, dtype=np.single)\n",
        "        return features, target_policy, target_value\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6oe6TjYZodpj"
      },
      "source": [
        "### AlphaZero Network"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "-FwMy49-od0M"
      },
      "outputs": [],
      "source": [
        "class AlphaZeroNetwork(nn.Module):\n",
        "    def __init__(self,\n",
        "                 input_channels: int,\n",
        "                 input_channel_height: int,\n",
        "                 input_channel_width: int,\n",
        "                 action_size: int,\n",
        "                 hidden_channels: int = 16,\n",
        "                 ):\n",
        "        super().__init__()\n",
        "        self.conv1 = nn.Conv2d(input_channels, hidden_channels, kernel_size=3, padding=1)\n",
        "        self.bn1 = nn.BatchNorm2d(hidden_channels)\n",
        "        self.conv2 = nn.Conv2d(hidden_channels, 1, kernel_size=1)\n",
        "        self.bn2 = nn.BatchNorm2d(1)\n",
        "        self.flat = nn.Flatten()\n",
        "        self.policy_head = nn.Linear(input_channel_height * input_channel_width, action_size)\n",
        "        self.value_head = nn.Linear(input_channel_height * input_channel_width, 1)\n",
        "\n",
        "    def forward(self, x):\n",
        "        x = nn.functional.relu(self.bn1(self.conv1(x)))\n",
        "        x = nn.functional.relu(self.bn2(self.conv2(x)))\n",
        "        x = self.flat(x)\n",
        "        policy_logit = self.policy_head(x)\n",
        "        policy = nn.functional.softmax(policy_logit, dim=1)\n",
        "        value = torch.tanh(self.value_head(x))\n",
        "        return policy_logit, policy, value\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ag14ObvZo6H0"
      },
      "source": [
        "### Logger"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "oRziFBR8o6aw"
      },
      "outputs": [],
      "source": [
        "def now_time() -> str:\n",
        "    return datetime.now().strftime(\"[%Y-%m-%d %H:%M:%S] \")\n",
        "\n",
        "\n",
        "class Logger:\n",
        "    def __init__(self, env: TicTacToe, device: torch.device | str = 'cpu'):\n",
        "        self.folder_name: str = f'{type(env).__name__}'\n",
        "        self.model_folder_name: str = f'{type(env).__name__}/model'\n",
        "        self.device = device\n",
        "        if os.path.exists(self.model_folder_name):\n",
        "            shutil.rmtree(self.model_folder_name)\n",
        "        os.makedirs(self.model_folder_name)\n",
        "\n",
        "\n",
        "    def write_log(self, write_message: str, timestamp: bool = True) -> None:\n",
        "        with open(f'{self.folder_name}/training_log.txt', 'a') as f:\n",
        "            if timestamp:\n",
        "                write_message = now_time() + write_message\n",
        "            f.write(write_message + '\\n')\n",
        "            print(write_message)\n",
        "\n",
        "    def save_network(self, network: nn.Module, iteration: int) -> None:\n",
        "        torch.jit.script(network).save(f'{self.model_folder_name}/weight_iter_{iteration}.pt')\n",
        "\n",
        "    def load_network(self, iteration: int) -> nn.Module:\n",
        "        return torch.jit.load(f'{self.model_folder_name}/weight_iter_{iteration}.pt',\n",
        "                              map_location=torch.device(self.device))\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "s2gjd3aQpgX4"
      },
      "source": [
        "## AlphaZero Algorithm"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6MCtN9QUqLNd"
      },
      "source": [
        "### Self-Play"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "1gI1SIooqJuK"
      },
      "outputs": [],
      "source": [
        "def self_play(env: TicTacToe,\n",
        "              network: nn.Module,\n",
        "              logger: Logger,\n",
        "              sp_num_games_per_iteration: int,\n",
        "              sp_mcts_simulation: int,\n",
        "              display_step: int = 10) -> list[Game]:\n",
        "    network.eval()\n",
        "    games = []\n",
        "    with torch.no_grad():\n",
        "        for i in range(1, sp_num_games_per_iteration + 1):\n",
        "            search_statistics = []\n",
        "            env.reset()\n",
        "            while not env.is_terminal():\n",
        "                mcts = MCTS(copy.deepcopy(env), network)\n",
        "                mcts.simulate(sp_mcts_simulation)\n",
        "                action = mcts.decide_action()\n",
        "                env.act(action)\n",
        "                search_statistics.append(mcts.get_normalize_child_visits())\n",
        "\n",
        "            assert len(env.action_history) == len(search_statistics)\n",
        "            game = Game(env.action_history[:], env.get_eval_score(), search_statistics)\n",
        "            games.append(game)\n",
        "\n",
        "            if i % display_step == 0:\n",
        "                logger.write_log(f'sp games: {i} / {sp_num_games_per_iteration}')\n",
        "    return games\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "symzwpM0qlnT"
      },
      "source": [
        "### Optimization"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "hD7SK1adqlvP"
      },
      "outputs": [],
      "source": [
        "def optimization(env: TicTacToe,\n",
        "                 network: nn.Module,\n",
        "                 logger: Logger,\n",
        "                 replay_buffer: ReplayBuffer,\n",
        "                 op_batch_size: int,\n",
        "                 op_training_steps: int,\n",
        "                 display_step: int = 10,\n",
        "                 learning_rate: float = 0.02,\n",
        "                 momentum: float = 0.9,\n",
        "                 weight_decay: float = 1e-4) -> None:\n",
        "    network.train()\n",
        "    optimizer = torch.optim.SGD(network.parameters(),\n",
        "                                lr=learning_rate,\n",
        "                                momentum=momentum,\n",
        "                                weight_decay=weight_decay)\n",
        "    for i in range(1, op_training_steps + 1):\n",
        "        optimizer.zero_grad()\n",
        "        features, target_policy, target_value = replay_buffer.sample_batch(copy.deepcopy(env), op_batch_size)\n",
        "        features_tensor = torch.from_numpy(features).to(DEVICE)\n",
        "        target_policy_tensor = torch.from_numpy(target_policy).to(DEVICE)\n",
        "        target_value_tensor = torch.from_numpy(target_value).unsqueeze(-1).to(DEVICE)\n",
        "\n",
        "        policy_logit, _, value = network(features_tensor)\n",
        "        log_policy = nn.functional.log_softmax(policy_logit, dim=1)\n",
        "        loss_policy = nn.functional.cross_entropy(log_policy, target_policy_tensor)\n",
        "        loss_value = nn.functional.mse_loss(value, target_value_tensor)\n",
        "        loss = loss_policy + loss_value\n",
        "        loss.backward()\n",
        "        optimizer.step()\n",
        "\n",
        "        if i % display_step == 0:\n",
        "            logger.write_log(f'op training steps: {i} / {op_training_steps}')\n",
        "            logger.write_log(f'\\tloss_policy: {loss_policy.item():.4f}', timestamp=False)\n",
        "            logger.write_log(f'\\tloss_value : {loss_value.item():.4f}', timestamp=False)\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Ez5WgT3QqvFR"
      },
      "source": [
        "### Evaluation"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9qnKzdB7qPxN"
      },
      "outputs": [],
      "source": [
        "def evaluation(env: TicTacToe,\n",
        "               logger: Logger | None,\n",
        "               iteration_a: int,\n",
        "               iteration_b: int,\n",
        "               num_games: int,\n",
        "               mcts_simulation: int,\n",
        "               use_softmax: bool = True,\n",
        "               grading: bool = False,\n",
        "               your_model = None):\n",
        "    if grading:\n",
        "        iteration_a = -1\n",
        "        iteration_b = 0\n",
        "        network_a = None\n",
        "        network_b = your_model\n",
        "        network_b.eval()\n",
        "    else:\n",
        "        network_a = logger.load_network(iteration_a)\n",
        "        network_b = logger.load_network(iteration_b)\n",
        "        network_a.eval()\n",
        "        network_b.eval()\n",
        "    winner_count = [{'iter': iteration_a, 'winP1': 0, 'winP2': 0, 'draw': 0, 'lossP1': 0, 'lossP2': 0, 'scores': []},\n",
        "                    {'iter': iteration_b, 'winP1': 0, 'winP2': 0, 'draw': 0, 'lossP1': 0, 'lossP2': 0, 'scores': []}]\n",
        "    with torch.no_grad():\n",
        "        for i in range(num_games):\n",
        "            env.reset()\n",
        "            while not env.is_terminal():\n",
        "                network = network_a if env.current_player == 1. else network_b\n",
        "                if network is None:\n",
        "                    action = np.random.choice(env.get_legal_actions())\n",
        "                else:\n",
        "                    mcts = MCTS(copy.deepcopy(env), network)\n",
        "                    mcts.simulate(mcts_simulation, use_dirichlet_noise=False)\n",
        "                    action = mcts.decide_action(use_softmax=use_softmax)\n",
        "                env.act(action)\n",
        "            eval_score = env.get_eval_score()\n",
        "            winner_count[0]['scores'].append(eval_score)\n",
        "            winner_count[1]['scores'].append(-eval_score)\n",
        "            if eval_score == 1.:\n",
        "                winner_count[0]['winP1'] += 1\n",
        "                winner_count[1]['lossP2'] += 1\n",
        "            elif eval_score == -1.:\n",
        "                winner_count[0]['lossP1'] += 1\n",
        "                winner_count[1]['winP2'] += 1\n",
        "            else:\n",
        "                winner_count[0]['draw'] += 1\n",
        "                winner_count[1]['draw'] += 1\n",
        "            # Swap player\n",
        "            network_a, network_b = network_b, network_a\n",
        "            winner_count[0], winner_count[1] = winner_count[1], winner_count[0]\n",
        "\n",
        "    sorted(winner_count, key=lambda item: item['iter'])\n",
        "    own_info = winner_count.pop()\n",
        "    opponent_info = winner_count.pop()\n",
        "    title_name = f'iteration {own_info.pop(\"iter\")} vs. {opponent_info.pop(\"iter\")}'\n",
        "    mean_score = sum(own_info.pop('scores')) / num_games\n",
        "    if grading:\n",
        "        print(f'your model vs. random selection agent: {own_info}')\n",
        "        print(f'your model vs. random selection agent win rate: {(mean_score + 1) / 2:.2%}')\n",
        "    else:\n",
        "        logger.write_log(f'\\t{title_name}: {own_info}', timestamp=False)\n",
        "        logger.write_log(f'\\t{title_name} win rate: {(mean_score + 1) / 2:.2%}', timestamp=False)\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "s0Op9M2_sH_7"
      },
      "source": [
        "## Training Model"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "kBFU2_wpsXov"
      },
      "outputs": [],
      "source": [
        "# hyperparameters\n",
        "max_iteration = 50\n",
        "sp_num_games_per_iteration = 10\n",
        "sp_mcts_simulation = 8\n",
        "buffer_num_iteration = 5\n",
        "op_batch_size = 32\n",
        "op_training_steps = 50\n",
        "eval_interval_iteration = 5\n",
        "eval_num_games = 100\n",
        "network_num_hidden_channels = 8\n",
        "\n",
        "random_seed = 600\n",
        "np.random.seed(random_seed)\n",
        "torch.manual_seed(random_seed)\n",
        "\n",
        "env = TicTacToe()\n",
        "# env = ConnectFour()\n",
        "network = AlphaZeroNetwork(input_channels=env.get_num_input_channels(),\n",
        "                           input_channel_height=env.get_input_channel_height(),\n",
        "                           input_channel_width=env.get_input_channel_width(),\n",
        "                           action_size=env.get_policy_size(),\n",
        "                           hidden_channels=network_num_hidden_channels\n",
        "                           ).to(DEVICE)\n",
        "replay_buffer = ReplayBuffer(buffer_num_iteration * sp_num_games_per_iteration)\n",
        "logger = Logger(env, device=DEVICE)\n",
        "logger.save_network(network, 0)\n",
        "\n",
        "for i in range(1, max_iteration + 1):\n",
        "    logger.write_log(f'iteration {i}:')\n",
        "\n",
        "    # Self-Play\n",
        "    latest_network = logger.load_network(i - 1)\n",
        "    games = self_play(env, latest_network, logger,\n",
        "                      sp_num_games_per_iteration=sp_num_games_per_iteration,\n",
        "                      sp_mcts_simulation=sp_mcts_simulation,\n",
        "                      display_step=10)\n",
        "    for game in games:\n",
        "        replay_buffer.save_game(game)\n",
        "\n",
        "    # Optimization\n",
        "    optimization(env, network, logger, replay_buffer,\n",
        "                 op_batch_size=op_batch_size,\n",
        "                 op_training_steps=op_training_steps,\n",
        "                 display_step=10)\n",
        "    logger.save_network(network, i)\n",
        "\n",
        "    # Evaluation\n",
        "    if i % eval_interval_iteration == 0:\n",
        "        logger.write_log('evaluation:')\n",
        "        evaluation(env, logger,\n",
        "                   iteration_a=0,\n",
        "                   iteration_b=i,\n",
        "                   num_games=eval_num_games,\n",
        "                   mcts_simulation=sp_mcts_simulation)\n",
        "        if i - eval_interval_iteration > 0:\n",
        "            evaluation(env, logger,\n",
        "                    iteration_a=i - eval_interval_iteration,\n",
        "                    iteration_b=i,\n",
        "                    num_games=eval_num_games,\n",
        "                    mcts_simulation=sp_mcts_simulation)\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "p95NO5GysEfb"
      },
      "source": [
        "## Evaluation (fight against the random selection agent)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "DwWoWsjfsFAy"
      },
      "outputs": [],
      "source": [
        "from google.colab import files\n",
        "\n",
        "random_seed = 600\n",
        "np.random.seed(random_seed)\n",
        "torch.manual_seed(random_seed)\n",
        "\n",
        "env = TicTacToe()\n",
        "# env = ConnectFour()\n",
        "\n",
        "print('env name:', type(env).__name__)\n",
        "\n",
        "# type your model file name\n",
        "fn = \"/content/TicTacToe/model/weight_iter_20.pt\"\n",
        "your_model = torch.jit.load(fn, map_location=torch.device(DEVICE))\n",
        "\n",
        "# or uploaded your model file (uncomment the code)\n",
        "\"\"\"\n",
        "uploaded = files.upload()\n",
        "if len(uploaded) != 1:\n",
        "    raise Exception('Only one file can be uploaded')\n",
        "your_model = None\n",
        "for fn in uploaded.keys():\n",
        "    your_model = torch.jit.load(fn, map_location=torch.device(DEVICE))\n",
        "\"\"\"\n",
        "\n",
        "evaluation(env, None,\n",
        "           iteration_a=1,\n",
        "           iteration_b=0,\n",
        "           num_games=100,\n",
        "           mcts_simulation=50,\n",
        "           use_softmax=False,\n",
        "           grading=True,\n",
        "           your_model=your_model)"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}