Multi Agent RL.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": false,
    "editable": false,
    "run_control": {
     "frozen": true
    }
   },
   "source": [
    "# Multi Agent Reinforcement Learning"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "deletable": false,
    "editable": false
   },
   "outputs": [],
   "source": [
    "from IPython.lib.display import YouTubeVideo"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": false,
    "editable": false,
    "run_control": {
     "frozen": true
    }
   },
   "source": [
    "### Markov Games"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "deletable": false,
    "editable": false
   },
   "outputs": [],
   "source": [
    "YouTubeVideo('Y9qq4Jqnwls')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": false,
    "editable": false,
    "run_control": {
     "frozen": true
    }
   },
   "source": [
    "### Approaches to Multi Agent Reinforcement Learning (MARL)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "deletable": false,
    "editable": false
   },
   "outputs": [],
   "source": [
    "YouTubeVideo('uKV9AJykin0')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": false,
    "editable": false,
    "run_control": {
     "frozen": true
    }
   },
   "source": [
    "### Cooperation, Competition, Mixed Environments"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "deletable": false,
    "editable": false
   },
   "outputs": [],
   "source": [
    "YouTubeVideo('vx6PIH5_oFg')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": false,
    "editable": false,
    "run_control": {
     "frozen": true
    }
   },
   "source": [
    "# Case Study: Physical Deception\n",
    "\n",
    "A paper on MARL, called **[Multi Agent Actor Critic for Mixed Cooperative Competitive environments](https://papers.nips.cc/paper/7217-multi-agent-actor-critic-for-mixed-cooperative-competitive-environments.pdf)** by *OpenAI* has used a team of 5 neural networks, called [OpenAI Five](https://openai.com/projects/five/) to defeat amateur [Dota 2](https://www.dota2.com/play/) players.\n",
    "\n",
    "We'll implement a part of this paper to train an agent to solve the **Physical Deception** problem."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "deletable": false,
    "editable": false
   },
   "outputs": [],
   "source": [
    "YouTubeVideo('nRKrQamUISs')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "deletable": false,
    "editable": false
   },
   "outputs": [],
   "source": [
    "YouTubeVideo('Ks9-TeCg3Fs')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "deletable": false,
    "editable": false
   },
   "outputs": [],
   "source": [
    "YouTubeVideo('4hFAhtLJR5U')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": false,
    "editable": false,
    "run_control": {
     "frozen": true
    }
   },
   "source": [
    "#### Objective of the Environment\n",
    "`Blue` dots are the **good agents**, and the `red` dot is an **adversary**. All of the agents' goals are to go near the `green` target. The blue agents know which one is green, but the red agent is color-blind and does not know which target is green/black! The optimal solution is for the red agent to chase one of the blue agent, and for the blue agents to split up and go toward each of the target.\n",
    "\n",
    "#### Running on Multi-core CPU the workspace\n",
    "Use of GPU wouldn't impact the training time for this program, Instead, Multicore environments would be a better choice to increase the training speed."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": false,
    "editable": false,
    "run_control": {
     "frozen": true
    }
   },
   "source": [
    "### Import the Packages"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "deletable": false,
    "editable": false
   },
   "outputs": [],
   "source": [
    "import os\n",
    "import time\n",
    "import random\n",
    "import numpy as np\n",
    "from copy import copy\n",
    "from math import sqrt\n",
    "from collections import deque\n",
    "\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "from torch.optim import Adam\n",
    "import torch.nn.functional as F\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "from matplotlib import animation\n",
    "%matplotlib inline\n",
    "\n",
    "from utils.ParallelEnvironments import MultiAgentParallelEnv\n",
    "\n",
    "from IPython.display import display as Display\n",
    "from IPython.display import HTML\n",
    "from pyvirtualdisplay import Display as display\n",
    "display = display(visible=0, size=(1400, 900))\n",
    "display.start()\n",
    "\n",
    "is_ipython = 'inline' in plt.get_backend()\n",
    "if is_ipython:\n",
    "    from IPython import display\n",
    "\n",
    "plt.ion()\n",
    "\n",
    "device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')\n",
    "device = 'cpu'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "MULTI_AGENT_ENVIRONMENT = 'simple_adversary.py'  # https://github.com/openai/multiagent-particle-envs\n",
    "MODEL_FILE              = './models/maddpg-adversary-episode-{:4d}.pt'\n",
    "N_ENVIRONMENTS          = 4\n",
    "RANDOM_SEED             = 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "deletable": false,
    "editable": false
   },
   "outputs": [],
   "source": [
    "random.seed(RANDOM_SEED)\n",
    "torch.manual_seed(RANDOM_SEED)\n",
    "torch.set_num_threads(N_ENVIRONMENTS)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": false,
    "editable": false,
    "run_control": {
     "frozen": true
    }
   },
   "source": [
    "### Explore the Multi Agent Environment"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "deletable": false,
    "editable": false
   },
   "outputs": [],
   "source": [
    "envs = MultiAgentParallelEnv(MULTI_AGENT_ENVIRONMENT, N_ENVIRONMENTS)\n",
    "\n",
    "print('Number of agents: {}'.format(envs.n_agents))\n",
    "print('Types of agents: {}'.format(envs.agent_types))\n",
    "print('Observation space of each agent: {}'.format(envs.states))\n",
    "print('Action space of each agent: {}'.format(envs.actions))\n",
    "\n",
    "envs.close()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": false,
    "editable": false,
    "run_control": {
     "frozen": true
    }
   },
   "source": [
    "### The Actor-Critic Network"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "code_folding": [
     0,
     6,
     45
    ],
    "deletable": false,
    "editable": false
   },
   "outputs": [],
   "source": [
    "def hidden_init(layer):\n",
    "    fan_in = layer.weight.data.size()[0]\n",
    "    lim = 1.0 / np.sqrt(fan_in)\n",
    "    return -lim, lim\n",
    "\n",
    "\n",
    "class Actor(nn.Module):\n",
    "    '''Actor (Policy) Model'''\n",
    "    def __init__(self, state_size, action_size, fc1_units=16, fc2_units=8):\n",
    "        '''Initialize parameters and build model.\n",
    "        Params\n",
    "        ======\n",
    "            state_size (int): Dimension of each state\n",
    "            action_size (int): Dimension of each action\n",
    "            fc1_units (int): Number of nodes in first hidden layer\n",
    "            fc2_units (int): Number of nodes in second hidden layer\n",
    "        '''\n",
    "        super(Actor, self).__init__()\n",
    "        self.fc1 = nn.Linear(state_size, fc1_units)\n",
    "        self.fc2 = nn.Linear(fc1_units, fc2_units)\n",
    "        self.fc3 = nn.Linear(fc2_units, action_size)\n",
    "        # self.input_norm = nn.BatchNorm1d(input_dim)\n",
    "        # self.input_norm.weight.data.fill_(1)\n",
    "        # self.input_norm.bias.data.fill_(0)        \n",
    "        # self.reset_parameters()\n",
    "\n",
    "    def reset_parameters(self):\n",
    "        self.fc1.weight.data.uniform_(*hidden_init(self.fc1))\n",
    "        self.fc2.weight.data.uniform_(*hidden_init(self.fc2))\n",
    "        self.fc3.weight.data.uniform_(-1e-3, 1e-3)\n",
    "\n",
    "    def forward(self, state):\n",
    "        '''\n",
    "        Build an actor (policy) network that maps states -> actions\n",
    "        Return a vector of the force\n",
    "        ''' \n",
    "        x = F.relu(self.fc1(state)) # F.leaky_relu\n",
    "        x = F.relu(self.fc2(x))\n",
    "        x = self.fc3(x)\n",
    "        norm = torch.norm(x)\n",
    "        # x is a 2D vector (a force that is applied to the agent)\n",
    "        # we bound the norm of the vector to be between 0 and 10\n",
    "        return 10.0 * (torch.tanh(norm)) * x / norm if norm > 0 else 10 * x\n",
    "\n",
    "\n",
    "class Critic(nn.Module):\n",
    "    '''Critic (Value) Model'''\n",
    "    def __init__(self, state_size, action_size, fc1_units=32, fc2_units=16):\n",
    "        '''Initialize parameters and build model.\n",
    "        Params\n",
    "        ======\n",
    "            state_size (int): Dimension of each state for individual agents\n",
    "            action_size (int): Dimension of ALL actions for ALL the agents\n",
    "            fc1_units (int): Number of nodes in the first hidden layer\n",
    "            fc2_units (int): Number of nodes in the second hidden layer\n",
    "        '''\n",
    "        super(Critic, self).__init__()\n",
    "        self.fc1 = nn.Linear(state_size + action_size, fc1_units)\n",
    "        self.fc2 = nn.Linear(fc1_units, fc2_units)\n",
    "        self.fc3 = nn.Linear(fc2_units, 1)\n",
    "        # self.reset_parameters()\n",
    "\n",
    "    def reset_parameters(self):\n",
    "        self.fc1.weight.data.uniform_(*hidden_init(self.fc1))\n",
    "        self.fc2.weight.data.uniform_(*hidden_init(self.fc2))\n",
    "        self.fc3.weight.data.uniform_(-1e-3, 1e-3)\n",
    "\n",
    "    def forward(self, state_action):\n",
    "        '''Build a critic (value) network that maps (state, action) pairs -> Q-values'''\n",
    "        x = F.relu(self.fc1(state_action)) # F.leaky_relu\n",
    "        x = F.relu(self.fc2(x))\n",
    "        x = self.fc3(x)\n",
    "        return x"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": false,
    "editable": false,
    "run_control": {
     "frozen": true
    }
   },
   "source": [
    "### Exploration: Adding random noise to the continuous actions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "code_folding": [
     0
    ],
    "deletable": false,
    "editable": false
   },
   "outputs": [],
   "source": [
    "class OUNoise:\n",
    "    def __init__(self, action_size, scale=0.1, mu=0, theta=0.15, sigma=0.2):\n",
    "        self.action_size = action_size\n",
    "        self.scale = scale\n",
    "        self.mu = mu\n",
    "        self.theta = theta\n",
    "        self.sigma = sigma\n",
    "        self.state = np.ones(self.action_size) * self.mu\n",
    "        self.reset()\n",
    "\n",
    "    def reset(self):\n",
    "        self.state = np.ones(self.action_size) * self.mu\n",
    "\n",
    "    def sample(self):\n",
    "        x = self.state\n",
    "        dx = self.theta * (self.mu - x) + self.sigma * np.random.randn(len(x))\n",
    "        self.state = x + dx\n",
    "        return torch.tensor(self.state * self.scale).float()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": false,
    "editable": false,
    "run_control": {
     "frozen": true
    }
   },
   "source": [
    "### Experience Replay Buffer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "code_folding": [
     0
    ],
    "deletable": false,
    "editable": false
   },
   "outputs": [],
   "source": [
    "class ReplayBuffer:\n",
    "    def __init__(self, size):\n",
    "        self.size = size\n",
    "        self.deque = deque(maxlen=self.size)\n",
    "\n",
    "    def push(self,transition):\n",
    "        'push into the buffer'\n",
    "        for item in list(map(list, zip(*transition))):\n",
    "            self.deque.append(item)\n",
    "\n",
    "    def sample(self, batchsize):\n",
    "        'sample from the buffer'\n",
    "        samples = random.sample(self.deque, batchsize)\n",
    "        return list(map(list, zip(*samples)))\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.deque)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": false,
    "editable": false,
    "run_control": {
     "frozen": true
    }
   },
   "source": [
    "### Hyperparameters for Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "BATCH_SIZE = 1000                # minibatch size\n",
    "GAMMA = 0.95                     # discount factor\n",
    "TAU = 0.02                       # for soft update of target parameters\n",
    "\n",
    "LR_ACTOR = 1e-2                  # learning rate of the actor \n",
    "LR_CRITIC = LR_ACTOR             # learning rate of the critic [1, 10] times of LR_ACTOR\n",
    "WEIGHT_DECAY = 1e-5              # L2 weight decay [0, 0.0001]\n",
    "\n",
    "TRAIN_EVERY = 2 * N_ENVIRONMENTS # number of episodes before update the model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": false,
    "editable": false,
    "run_control": {
     "frozen": true
    }
   },
   "source": [
    "### The DDPG Agent"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "code_folding": [
     0
    ],
    "deletable": false,
    "editable": false
   },
   "outputs": [],
   "source": [
    "class DDPGAgent:\n",
    "    def __init__(self, state_size, action_size, n_agents):\n",
    "        '''\n",
    "        Params\n",
    "        ======\n",
    "            state_size (int): dimension of each state\n",
    "            action_size (int): dimension of each action\n",
    "            n_agents (int): number of agents in the environment\n",
    "        '''\n",
    "        super(DDPGAgent, self).__init__()\n",
    "        self.actor = Actor(state_size, action_size).to(device)\n",
    "        self.target_actor = Actor(state_size, action_size).to(device)\n",
    "        self.critic = Critic(state_size, action_size * n_agents).to(device)\n",
    "        self.target_critic = Critic(state_size, action_size * n_agents).to(device)\n",
    "        \n",
    "        self.noise = OUNoise(action_size, scale=1.0)\n",
    "\n",
    "        # initialize targets same as original networks\n",
    "        self.hard_update(self.target_actor, self.actor)\n",
    "        self.hard_update(self.target_critic, self.critic)\n",
    "\n",
    "        self.actor_optimizer = Adam(self.actor.parameters(), lr=LR_ACTOR)\n",
    "        self.critic_optimizer = Adam(self.critic.parameters(), lr=LR_CRITIC, weight_decay=WEIGHT_DECAY)\n",
    "\n",
    "    def act(self, obs, noise=0.0):\n",
    "        obs = obs.to(device)\n",
    "        action = self.actor(obs) + noise * self.noise.sample()\n",
    "        return action\n",
    "\n",
    "    def target_act(self, obs, noise=0.0):\n",
    "        obs = obs.to(device)\n",
    "        action = self.target_actor(obs) + noise * self.noise.sample()\n",
    "        return action\n",
    "    \n",
    "    def hard_update(self, target_model, source_model):\n",
    "        for target_param, source_param in zip(target_model.parameters(), source_model.parameters()):\n",
    "            target_param.data.copy_(source_param.data)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": false,
    "editable": false,
    "run_control": {
     "frozen": true
    }
   },
   "source": [
    "### Multi Agent DDPG"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "code_folding": [
     0,
     4
    ],
    "deletable": false,
    "editable": false
   },
   "outputs": [],
   "source": [
    "def transpose_to_tensor(input_list):\n",
    "    make_tensor = lambda x: torch.tensor(x, dtype=torch.float)\n",
    "    return list(map(make_tensor, zip(*input_list)))\n",
    "\n",
    "class MADDPG:\n",
    "    def __init__(self, state_size, action_size, n_agents):\n",
    "        super(MADDPG, self).__init__()\n",
    "        self.maddpg_agent = [DDPGAgent(state_size, action_size, n_agents) for _ in range(n_agents)]        \n",
    "        self.iter = 0\n",
    "\n",
    "    def get_actors(self):\n",
    "        '''get actors of all the agents in the MADDPG object'''\n",
    "        return [ddpg_agent.actor for ddpg_agent in self.maddpg_agent]\n",
    "\n",
    "    def get_target_actors(self):\n",
    "        '''get target_actors of all the agents in the MADDPG object'''\n",
    "        return [ddpg_agent.target_actor for ddpg_agent in self.maddpg_agent]\n",
    "\n",
    "    def act(self, obs_all_agents, noise=0.0):\n",
    "        '''get actions from all agents in the MADDPG object'''\n",
    "        actions = [agent.act(obs, noise) for agent, obs in zip(self.maddpg_agent, obs_all_agents)]\n",
    "        return actions\n",
    "\n",
    "    def target_act(self, obs_all_agents, noise=0.0):\n",
    "        '''get target network actions from all the agents in the MADDPG object'''\n",
    "        target_actions = [ddpg_agent.target_act(obs, noise) for ddpg_agent, obs in zip(self.maddpg_agent, obs_all_agents)]\n",
    "        return target_actions\n",
    "\n",
    "    def update(self, samples, agent_number):\n",
    "        '''\n",
    "        update the critic and actor of the agent referred by 'agent_number'\n",
    "        '''\n",
    "        \n",
    "        # need to transpose each element of the samples\n",
    "        # to flip obs[parallel_agent][agent_number] to\n",
    "        # obs[agent_number][parallel_agent]\n",
    "        obs, obs_full, action, reward, next_obs, next_obs_full, done = map(transpose_to_tensor, samples)\n",
    "\n",
    "        obs_full = torch.stack(obs_full)\n",
    "        next_obs_full = torch.stack(next_obs_full)\n",
    "        \n",
    "        agent = self.maddpg_agent[agent_number]\n",
    "        \n",
    "        # -------------- update critic network -------------\n",
    "        agent.critic_optimizer.zero_grad()\n",
    "\n",
    "        # critic loss = batch mean of (y- Q(s,a) from target network)^2\n",
    "        # y = reward of this timestep + discount * Q(s', a') from target network\n",
    "        target_actions = self.target_act(next_obs)\n",
    "        target_actions = torch.cat(target_actions, dim=1)\n",
    "        \n",
    "        target_critic_input = torch.cat((next_obs_full.t(), target_actions), dim=1).to(device)\n",
    "        \n",
    "        with torch.no_grad():\n",
    "            q_next = agent.target_critic(target_critic_input)\n",
    "        \n",
    "        y = reward[agent_number].view(-1, 1) + GAMMA * q_next * (1 - done[agent_number].view(-1, 1))\n",
    "        action = torch.cat(action, dim=1)\n",
    "        critic_input = torch.cat((obs_full.t(), action), dim=1).to(device)\n",
    "        q = agent.critic(critic_input)\n",
    "\n",
    "        huber_loss = torch.nn.SmoothL1Loss()\n",
    "        critic_loss = huber_loss(q, y.detach())\n",
    "        critic_loss.backward()\n",
    "        # torch.nn.utils.clip_grad_norm_(agent.critic.parameters(), 0.5)\n",
    "        agent.critic_optimizer.step()\n",
    "\n",
    "        # ------------- update actor network using policy gradient -------------\n",
    "        agent.actor_optimizer.zero_grad()\n",
    "        # make input to agent\n",
    "        # detach the other agents to save computation\n",
    "        # saves some time for computing derivative\n",
    "        q_input = [self.maddpg_agent[i].actor(ob) if i == agent_number\n",
    "                   else self.maddpg_agent[i].actor(ob).detach()\n",
    "                   for i, ob in enumerate(obs)]\n",
    "                \n",
    "        q_input = torch.cat(q_input, dim=1)\n",
    "        \n",
    "        # combine all the actions and observations for input to critic\n",
    "        # many of the obs are redundant, and obs[1] contains all useful information already\n",
    "        q_input = torch.cat((obs_full.t(), q_input), dim=1)\n",
    "        \n",
    "        # get the policy gradient\n",
    "        actor_loss = - agent.critic(q_input).mean()\n",
    "        actor_loss.backward()\n",
    "        # torch.nn.utils.clip_grad_norm_(agent.actor.parameters(),0.5)\n",
    "        agent.actor_optimizer.step()\n",
    "\n",
    "    def update_targets(self):\n",
    "        '''soft update targets'''\n",
    "        self.iter += 1\n",
    "        for ddpg_agent in self.maddpg_agent:\n",
    "            self.soft_update(ddpg_agent.target_actor, ddpg_agent.actor)\n",
    "            self.soft_update(ddpg_agent.target_critic, ddpg_agent.critic)\n",
    "    \n",
    "    def soft_update(self, target, source):\n",
    "        for target_param, param in zip(target.parameters(), source.parameters()):\n",
    "            target_param.data.copy_(target_param.data * (1.0 - TAU) + param.data * TAU)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": false,
    "editable": false,
    "run_control": {
     "frozen": true
    }
   },
   "source": [
    "### The Training Loop"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "code_folding": [
     0
    ],
    "deletable": false,
    "editable": false
   },
   "outputs": [],
   "source": [
    "def train(n_episodes=1000, planning_horizon=80, print_every=100, save_every=10):\n",
    "    \n",
    "    # amplitude of OU noise, this slowly decreases to 0\n",
    "    noise = 2\n",
    "    noise_reduction = 0.9999\n",
    "    \n",
    "    envs = MultiAgentParallelEnv(MULTI_AGENT_ENVIRONMENT, N_ENVIRONMENTS)\n",
    "    agent_types = envs.agent_types\n",
    "    n_agents = envs.n_agents\n",
    "    state_size = envs.states[0].shape[0]\n",
    "    action_size = envs.actions[0].shape[0]\n",
    "    \n",
    "    buffer = ReplayBuffer(5000 * planning_horizon) # keep 5000 episodes worth of replay\n",
    "    maddpg = MADDPG(state_size, action_size, n_agents) # initialize policy and critic\n",
    "    agent_reward = [[] for _ in range(n_agents)]\n",
    "\n",
    "    for episode in range(0, n_episodes, N_ENVIRONMENTS):\n",
    "        reward_this_episode = np.zeros((N_ENVIRONMENTS, 3))\n",
    "        obs, obs_full = envs.reset()\n",
    "        for episode_t in range(planning_horizon):\n",
    "            actions = maddpg.act(transpose_to_tensor(obs), noise=noise)\n",
    "            noise *= noise_reduction\n",
    "            actions_array = torch.stack(actions).detach().numpy()\n",
    "\n",
    "            # transpose the list of list\n",
    "            # flip the first two indices\n",
    "            # input to step requires the first index to correspond to number of parallel agents\n",
    "            actions_for_env = np.rollaxis(actions_array, 1)\n",
    "            \n",
    "            # step forward one frame\n",
    "            next_obs, next_obs_full, rewards, dones, info = envs.step(actions_for_env)\n",
    "            \n",
    "            # add data to buffer\n",
    "            transition = (obs, obs_full, actions_for_env, rewards, next_obs, next_obs_full, dones)\n",
    "            buffer.push(transition)\n",
    "            reward_this_episode += rewards\n",
    "            obs, obs_full = next_obs, next_obs_full\n",
    "\n",
    "        for env_idx in range(N_ENVIRONMENTS):\n",
    "            for agent_idx in range(n_agents):\n",
    "                agent_reward[agent_idx].append(reward_this_episode[env_idx, agent_idx])\n",
    "        \n",
    "        # update once after every TRAIN_EVERY episodes\n",
    "        if len(buffer) > BATCH_SIZE and episode % TRAIN_EVERY < N_ENVIRONMENTS:\n",
    "            for agent_idx in range(n_agents):\n",
    "                samples = buffer.sample(BATCH_SIZE)\n",
    "                maddpg.update(samples, agent_idx)\n",
    "            maddpg.update_targets() # soft update the target network towards the actual networks\n",
    "\n",
    "        # display progress after every episode        \n",
    "        print('\\r' + ' ' * 120, end='')\n",
    "        print('\\rEpisode: {:04d}/{:04d}\\tAvg. Rewards for {} Agents: {}'.format(episode, n_episodes, n_agents, [np.round(np.mean(agent_reward[i]), 4) for i in range(n_agents)]), end='')\n",
    "        if (episode + 1) % print_every < N_ENVIRONMENTS:\n",
    "            print('\\r' + ' ' * 120, end='')\n",
    "            print('\\rEpisode: {:04d}/{:04d}\\tAvg. Rewards for {} Agents: {}'.format(episode, n_episodes, n_agents, [np.round(np.mean(agent_reward[i]), 4) for i in range(n_agents)]))\n",
    "        if episode % save_every < N_ENVIRONMENTS or episode == n_episodes - N_ENVIRONMENTS:\n",
    "            torch.save([\n",
    "                {\n",
    "                    'actor_params' : maddpg.maddpg_agent[i].actor.state_dict(),\n",
    "                    'actor_optim_params': maddpg.maddpg_agent[i].actor_optimizer.state_dict(),\n",
    "                    'critic_params' : maddpg.maddpg_agent[i].critic.state_dict(),\n",
    "                    'critic_optim_params' : maddpg.maddpg_agent[i].critic_optimizer.state_dict()\n",
    "                } for i in range(n_agents)\n",
    "            ], MODEL_FILE.format(episode))\n",
    "\n",
    "    envs.close()\n",
    "    return agent_reward, agent_types"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "deletable": false,
    "editable": false,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "agent_rewards, agent_types = train()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": false,
    "editable": false,
    "run_control": {
     "frozen": true
    }
   },
   "source": [
    "### Plot the Scores"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "code_folding": [
     2
    ],
    "deletable": false,
    "editable": false
   },
   "outputs": [],
   "source": [
    "fig, axs = plt.subplots(3, 1, sharex=True, figsize=(10, 10))\n",
    "axs = axs.flatten()\n",
    "for i, (title, scores) in enumerate(zip(agent_types, agent_rewards)):\n",
    "    axs[i].plot(np.arange(1, len(scores)+1), scores)\n",
    "    axs[i].set_title(title)\n",
    "    axs[i].set_ylabel('Score')\n",
    "    axs[i].set_xlabel('' if i != len(agent_types) - 1 else 'Episode #')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": false,
    "editable": false,
    "run_control": {
     "frozen": true
    }
   },
   "source": [
    "### Watch the Smart Agents in Action"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "LOAD_EPISODE = 996\n",
    "LOAD_N_ENVIRONMENTS = 4\n",
    "SHOW_CONTROL = True"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "code_folding": [
     0,
     18,
     22
    ],
    "deletable": false,
    "editable": false,
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "def animate_frames(frames):\n",
    "    'function to animate a list of frames'\n",
    "    def display_animation(anim):\n",
    "        plt.close(anim._fig)\n",
    "        return HTML(anim.to_jshtml())\n",
    "    plt.axis('off')\n",
    "    cmap = None if len(frames[0].shape) == 3 else 'Greys' # color option for plotting, use Greys for greyscale\n",
    "    patch = plt.imshow(frames[0], cmap=cmap)  \n",
    "    fanim = animation.FuncAnimation(plt.gcf(), lambda x: patch.set_data(frames[x]), frames = len(frames), interval=30)\n",
    "    Display(display_animation(fanim))\n",
    "\n",
    "envs = MultiAgentParallelEnv(MULTI_AGENT_ENVIRONMENT, LOAD_N_ENVIRONMENTS) # seed=int(time.time())\n",
    "n_agents = envs.n_agents\n",
    "maddpg = MADDPG(envs.states[0].shape[0], envs.actions[0].shape[0], n_agents)  \n",
    "\n",
    "# load the weights from file\n",
    "map_location = (lambda storage, loc: storage.cuda()) if torch.cuda.is_available() else 'cpu'\n",
    "checkpoint = torch.load(MODEL_FILE.format(LOAD_EPISODE), map_location=map_location)        \n",
    "for i in range(n_agents):\n",
    "    maddpg.maddpg_agent[i].actor.load_state_dict(checkpoint[i]['actor_params'])\n",
    "    maddpg.maddpg_agent[i].actor.eval()\n",
    "\n",
    "for e in range(1):\n",
    "    frames = []\n",
    "    obs, _ = envs.reset()\n",
    "    if SHOW_CONTROL:\n",
    "        frames.append(envs.render(mode='rgb_array'))\n",
    "    else:\n",
    "        img = plt.imshow(envs.render(mode='rgb_array'))\n",
    "    for t in range(200):\n",
    "        actions = maddpg.act(transpose_to_tensor(obs), noise=0)\n",
    "        actions_array = torch.stack(actions).detach().numpy()\n",
    "        actions_for_env = np.rollaxis(actions_array, 1)\n",
    "        if SHOW_CONTROL:\n",
    "            frames.append(envs.render(mode='rgb_array'))\n",
    "        else:\n",
    "            img.set_data(envs.render(mode='rgb_array')) \n",
    "            plt.axis('off')\n",
    "            display.display(plt.gcf())\n",
    "            display.clear_output(wait=True)\n",
    "        obs, _, _, dones, _ = envs.step(actions_for_env)\n",
    "        if dones.any():\n",
    "            break\n",
    "    if SHOW_CONTROL:\n",
    "        animate_frames(frames)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": false,
    "editable": false,
    "run_control": {
     "frozen": true
    }
   },
   "source": [
    "# Case Study: AlphaZero\n",
    "\n",
    "The following materials are derived from the original papers, [alphago zero](https://deepmind.com/documents/119/agz_unformatted_nature.pdf), and [alphazero](https://arxiv.org/abs/1712.01815) by the DeepMind."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "deletable": false,
    "editable": false
   },
   "outputs": [],
   "source": [
    "YouTubeVideo('Zzc1XJ1aJ-4')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": false,
    "editable": false,
    "run_control": {
     "frozen": true
    }
   },
   "source": [
    "## Zero-Sum Game"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "deletable": false,
    "editable": false
   },
   "outputs": [],
   "source": [
    "YouTubeVideo('uPw1dHVqdXQ')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": false,
    "editable": false,
    "run_control": {
     "frozen": true
    }
   },
   "source": [
    "## Monte Carlo Tree Search\n",
    "\n",
    "### Random Sampling"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "deletable": false,
    "editable": false
   },
   "outputs": [],
   "source": [
    "YouTubeVideo('wn2B3j_Qz6E')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": false,
    "editable": false,
    "run_control": {
     "frozen": true
    }
   },
   "source": [
    "### Expansion and Back-propagation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "deletable": false,
    "editable": false
   },
   "outputs": [],
   "source": [
    "YouTubeVideo('H34Wtk1iNDY')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": false,
    "editable": false,
    "run_control": {
     "frozen": true
    }
   },
   "source": [
    "## AlphaZero\n",
    "\n",
    "### Guided Tree Search"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "deletable": false,
    "editable": false
   },
   "outputs": [],
   "source": [
    "YouTubeVideo('LinuRy47xbw')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": false,
    "editable": false,
    "run_control": {
     "frozen": true
    }
   },
   "source": [
    "### Self-Play Training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "deletable": false,
    "editable": false
   },
   "outputs": [],
   "source": [
    "YouTubeVideo('wl1qfPXqRuQ')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": false,
    "editable": false,
    "run_control": {
     "frozen": true
    }
   },
   "source": [
    "---\n",
    "\n",
    "# TicTacToe using AlphaZero"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": false,
    "editable": false,
    "run_control": {
     "frozen": true
    }
   },
   "source": [
    "We'll now implement an advanced version of `TicTacToe` using AlphaZero.\n",
    "\n",
    "For code walkthroughs, you can watch these [notebook walkthrough](https://www.youtube.com/watch?v=uUFuBscf98I) and [python classes walkthrough](https://www.youtube.com/watch?v=hKnBQvtJ_zQ) videos."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": false,
    "editable": false,
    "run_control": {
     "frozen": true
    }
   },
   "source": [
    "### Implementation of the Monte Carlo Tree Search (MCTS) Algorithm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "code_folding": [
     15,
     37,
     56
    ],
    "deletable": false,
    "editable": false
   },
   "outputs": [],
   "source": [
    "c = 1.0\n",
    "\n",
    "# transformations\n",
    "t0= lambda x: x\n",
    "t1= lambda x: x[:,::-1].copy()\n",
    "t2= lambda x: x[::-1,:].copy()\n",
    "t3= lambda x: x[::-1,::-1].copy()\n",
    "t4= lambda x: x.T\n",
    "t5= lambda x: x[:,::-1].T.copy()\n",
    "t6= lambda x: x[::-1,:].T.copy()\n",
    "t7= lambda x: x[::-1,::-1].T.copy()\n",
    "\n",
    "tlist=[t0, t1,t2,t3,t4,t5,t6,t7]\n",
    "tlist_half=[t0,t1,t2,t3]\n",
    "\n",
    "def flip(x, dim):\n",
    "    indices = [slice(None)] * x.dim()\n",
    "    indices[dim] = torch.arange(x.size(dim) - 1, -1, -1, dtype=torch.long, device=x.device)\n",
    "    return x[tuple(indices)]\n",
    "\n",
    "\n",
    "t0inv= lambda x: x\n",
    "t1inv= lambda x: flip(x,1)\n",
    "t2inv= lambda x: flip(x,0)\n",
    "t3inv= lambda x: flip(flip(x,0),1)\n",
    "t4inv= lambda x: x.t()\n",
    "t5inv= lambda x: flip(x,0).t()\n",
    "t6inv= lambda x: flip(x,1).t()\n",
    "t7inv= lambda x: flip(flip(x,0),1).t()\n",
    "\n",
    "tinvlist = [t0inv, t1inv, t2inv, t3inv, t4inv, t5inv, t6inv, t7inv]\n",
    "tinvlist_half=[t0inv, t1inv, t2inv, t3inv]\n",
    "\n",
    "transformation_list = list(zip(tlist, tinvlist))\n",
    "transformation_list_half = list(zip(tlist_half, tinvlist_half))\n",
    "\n",
    "\n",
    "def process_policy(policy, game):\n",
    "\n",
    "    # for square board, add rotations as well\n",
    "    if game.size[0]==game.size[1]:\n",
    "        t, tinv = random.choice(transformation_list)\n",
    "\n",
    "    # otherwise only add reflections\n",
    "    else:\n",
    "        t, tinv = random.choice(transformation_list_half)\n",
    "     \n",
    "    frame=torch.tensor(t(game.state*game.player), dtype=torch.float, device=device)\n",
    "    input=frame.unsqueeze(0).unsqueeze(0)\n",
    "    prob, v = policy(input)\n",
    "    mask = torch.tensor(game.available_mask(), dtype=torch.bool)\n",
    "    \n",
    "    # we add a negative sign because when deciding next move,\n",
    "    # the current player is the previous player making the move\n",
    "    return game.available_moves(), tinv(prob)[mask].view(-1), v.squeeze().squeeze()\n",
    "\n",
    "class MCTSNode:\n",
    "    def __init__(self, game, mother=None, prob=torch.tensor(0., dtype=torch.float)):\n",
    "        self.game = game\n",
    "          \n",
    "        # child nodes\n",
    "        self.child = {}\n",
    "        # numbers for determining which actions to take next\n",
    "        self.U = 0\n",
    "\n",
    "        # V from neural net output\n",
    "        # it's a torch.tensor object\n",
    "        # has require_grad enabled\n",
    "        self.prob = prob\n",
    "        # the predicted expectation from neural net\n",
    "        self.nn_v = torch.tensor(0., dtype=torch.float)\n",
    "        \n",
    "        # visit count\n",
    "        self.N = 0\n",
    "\n",
    "        # expected V from MCTS\n",
    "        self.V = 0\n",
    "\n",
    "        # keeps track of the guaranteed outcome\n",
    "        # initialized to None\n",
    "        # this is for speeding the tree-search up\n",
    "        # but stopping exploration when the outcome is certain\n",
    "        # and there is a known perfect play\n",
    "        self.outcome = self.game.score\n",
    "\n",
    "\n",
    "        # if game is won/loss/draw\n",
    "        if self.game.score is not None:\n",
    "            self.V = self.game.score*self.game.player\n",
    "            self.U = 0 if self.game.score is 0 else self.V*float('inf')\n",
    "\n",
    "        # link to previous node\n",
    "        self.mother = mother\n",
    "\n",
    "    def create_child(self, actions, probs):\n",
    "        # create a dictionary of children\n",
    "        games = [ copy(self.game) for a in actions ]\n",
    "\n",
    "        for action, game in zip(actions, games):\n",
    "            game.move(action)\n",
    "\n",
    "        child = { tuple(a): MCTSNode(g, self, p) for a,g,p in zip(actions, games, probs) }\n",
    "        self.child = child\n",
    "        \n",
    "    def explore(self, policy):\n",
    "\n",
    "        if self.game.score is not None:\n",
    "            raise ValueError(\"game has ended with score {0:d}\".format(self.game.score))\n",
    "\n",
    "        current = self\n",
    "\n",
    "        \n",
    "        # explore children of the node\n",
    "        # to speed things up \n",
    "        while current.child and current.outcome is None:\n",
    "\n",
    "            child = current.child\n",
    "            max_U = max(c.U for c in child.values())\n",
    "            # print(\"current max_U \", max_U) \n",
    "            actions = [ a for a,c in child.items() if c.U == max_U ]\n",
    "            if len(actions) == 0:\n",
    "                print(\"error zero length \", max_U)\n",
    "                print(current.game.state)\n",
    "                      \n",
    "            action = random.choice(actions)            \n",
    "\n",
    "            if max_U == -float(\"inf\"):\n",
    "                current.U = float(\"inf\")\n",
    "                current.V = 1.0\n",
    "                break\n",
    "            \n",
    "            elif max_U == float(\"inf\"):\n",
    "                current.U = -float(\"inf\")\n",
    "                current.V = -1.0\n",
    "                break\n",
    "                \n",
    "            current = child[action]\n",
    "        \n",
    "        # if node hasn't been expanded\n",
    "        if not current.child and current.outcome is None:\n",
    "            # policy outputs results from the perspective of the next player\n",
    "            # thus extra - sign is needed\n",
    "            next_actions, probs, v = process_policy(policy, current.game)\n",
    "            current.nn_v = -v\n",
    "            current.create_child(next_actions, probs)\n",
    "            current.V = -float(v)\n",
    "\n",
    "        \n",
    "        current.N += 1\n",
    "\n",
    "        # now update U and back-prop\n",
    "        while current.mother:\n",
    "            mother = current.mother\n",
    "            mother.N += 1\n",
    "            # beteen mother and child, the player is switched, extra - sign\n",
    "            mother.V += (-current.V - mother.V)/mother.N\n",
    "\n",
    "            #update U for all sibling nodes\n",
    "            for sibling in mother.child.values():\n",
    "                if sibling.U is not float(\"inf\") and sibling.U is not -float(\"inf\"):\n",
    "                    sibling.U = sibling.V + c*float(sibling.prob)* sqrt(mother.N)/(1+sibling.N)\n",
    "\n",
    "            current = current.mother\n",
    "\n",
    "\n",
    "               \n",
    "    def next(self, temperature=1.0):\n",
    "\n",
    "        if self.game.score is not None:\n",
    "            raise ValueError('game has ended with score {0:d}'.format(self.game.score))\n",
    "\n",
    "        if not self.child:\n",
    "            print(self.game.state)\n",
    "            raise ValueError('no children found and game hasn\\'t ended')\n",
    "        \n",
    "        child=self.child\n",
    "\n",
    "        \n",
    "        # if there are winning moves, just output those\n",
    "        max_U = max(c.U for c in child.values())\n",
    "\n",
    "        if max_U == float(\"inf\"):\n",
    "            prob = torch.tensor([ 1.0 if c.U == float(\"inf\") else 0 for c in child.values()], device=device)\n",
    "            \n",
    "        else:\n",
    "            # divide things by maxN for numerical stability\n",
    "            maxN = max(node.N for node in child.values())+1\n",
    "            prob = torch.tensor([ (node.N/maxN)**(1/temperature) for node in child.values() ], device=device)\n",
    "\n",
    "        # normalize the probability\n",
    "        if torch.sum(prob) > 0:\n",
    "            prob /= torch.sum(prob)\n",
    "            \n",
    "        # if sum is zero, just make things random\n",
    "        else:\n",
    "            prob = torch.tensor(1.0/len(child), device=device).repeat(len(child))\n",
    "\n",
    "        nn_prob = torch.stack([ node.prob for node in child.values() ]).to(device)\n",
    "\n",
    "        nextstate = random.choices(list(child.values()), weights=prob)[0]\n",
    "        \n",
    "        # V was for the previous player making a move\n",
    "        # to convert to the current player we add - sign\n",
    "        return nextstate, (-self.V, -self.nn_v, prob, nn_prob)\n",
    "\n",
    "    def detach_mother(self):\n",
    "        del self.mother\n",
    "        self.mother = None"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": false,
    "editable": false,
    "run_control": {
     "frozen": true
    }
   },
   "source": [
    "### Utility Methods and Classes to Set Up a Game"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "code_folding": [
     2,
     10,
     11,
     17,
     57
    ],
    "deletable": false,
    "editable": false
   },
   "outputs": [],
   "source": [
    "# output the index of when v has a continuous string of i\n",
    "# get_runs([0,0,1,1,1,0,0],1) gives [2],[5],[3]\n",
    "def get_runs(v, i):\n",
    "     bounded = np.hstack(([0], (v==i).astype(int), [0]))\n",
    "     difs = np.diff(bounded)\n",
    "     starts, = np.where(difs > 0)\n",
    "     ends, = np.where(difs < 0)\n",
    "     return starts, ends, ends-starts\n",
    "\n",
    "# see if vector contains N of certain number in a row\n",
    "def in_a_row(v, N, i):\n",
    "     if len(v) < N:\n",
    "          return False\n",
    "     else:\n",
    "          _, _, total = get_runs(v,i)\n",
    "          return np.any(total >= N)\n",
    "        \n",
    "def get_lines(matrix, loc):\n",
    "\n",
    "     i,j=loc\n",
    "     flat = matrix.reshape(-1,*matrix.shape[2:])\n",
    "    \n",
    "     w = matrix.shape[0]\n",
    "     h = matrix.shape[1]\n",
    "     def flat_pos(pos):\n",
    "          return pos[0]*h+pos[1]\n",
    "\n",
    "     pos = flat_pos((i,j))\n",
    "\n",
    "     # index for flipping matrix across different axis\n",
    "     ic = w-1-i\n",
    "     jc = h-1-j\n",
    "\n",
    "     # top left\n",
    "     tl = (i-j,0) if i>j else (0, j-i)\n",
    "     tl = flat_pos(tl)\n",
    "\n",
    "     # bottom left\n",
    "     bl = (w-1-(ic-j),0) if ic>j else (w-1, j-ic)\n",
    "     bl = flat_pos(bl)\n",
    "\n",
    "     # top right\n",
    "     tr = (i-jc,h-1) if i>jc else (0, h-1-(jc-i))\n",
    "     tr = flat_pos(tr)\n",
    "\n",
    "     # bottom right\n",
    "     br = (w-1-(ic-jc),h-1) if ic>jc else (w-1, h-1-(jc-ic))\n",
    "     br = flat_pos(br)\n",
    "\n",
    "     hor = matrix[:,j]\n",
    "     ver = matrix[i,:]\n",
    "     diag_right = np.concatenate([flat[tl:pos:h+1],flat[pos:br+1:h+1]])\n",
    "     diag_left = np.concatenate([flat[tr:pos:h-1],flat[pos:bl+1:h-1]])\n",
    "\n",
    "     return hor, ver, diag_right, diag_left\n",
    "\n",
    "\n",
    "class ConnectN:\n",
    "\n",
    "     def __init__(self, size, N, pie_rule=False):\n",
    "          self.size = size\n",
    "          self.w, self.h = size\n",
    "          self.N = N\n",
    "\n",
    "          # make sure game is well defined\n",
    "          if self.w<0 or self.h<0 or self.N<2 or \\\n",
    "             (self.N > self.w and self.N > self.h):\n",
    "               raise ValueError('Game cannot initialize with a {0:d}x{1:d} grid, and winning condition {2:d} in a row'.format(self.w, self.h, self.N))\n",
    "\n",
    "          \n",
    "          self.score = None\n",
    "          self.state=np.zeros(size, dtype=np.float)\n",
    "          self.player=1\n",
    "          self.last_move=None\n",
    "          self.n_moves=0\n",
    "          self.pie_rule=pie_rule\n",
    "          self.switched_side=False\n",
    "\n",
    "     # fast deepcopy\n",
    "     def __copy__(self):\n",
    "          cls = self.__class__\n",
    "          new_game = cls.__new__(cls)\n",
    "          new_game.__dict__.update(self.__dict__)\n",
    "\n",
    "          new_game.N = self.N\n",
    "          new_game.pie_rule = self.pie_rule\n",
    "          new_game.state = self.state.copy()\n",
    "          new_game.switched_side = self.switched_side\n",
    "          new_game.n_moves = self.n_moves\n",
    "          new_game.last_move = self.last_move\n",
    "          new_game.player = self.player\n",
    "          new_game.score = self.score\n",
    "          return new_game\n",
    "    \n",
    "     # check victory condition\n",
    "     # fast version\n",
    "     def get_score(self):\n",
    "\n",
    "          # game cannot end beca\n",
    "          if self.n_moves<2*self.N-1:\n",
    "               return None\n",
    "\n",
    "          i,j = self.last_move\n",
    "          hor, ver, diag_right, diag_left = get_lines(self.state, (i,j))\n",
    "\n",
    "          # loop over each possibility\n",
    "          for line in [ver, hor, diag_right, diag_left]:\n",
    "               if in_a_row(line, self.N, self.player):\n",
    "                    return self.player\n",
    "                    \n",
    "          # no more moves\n",
    "          if np.all(self.state!=0):\n",
    "               return 0\n",
    "\n",
    "          return None\n",
    "\n",
    "     # for rendering\n",
    "     # output a list of location for the winning line\n",
    "     def get_winning_loc(self):\n",
    "        \n",
    "          if self.n_moves<2*self.N-1:\n",
    "               return []\n",
    "\n",
    "          \n",
    "          loc = self.last_move\n",
    "          hor, ver, diag_right, diag_left = get_lines(self.state, loc)\n",
    "          ind = np.indices(self.state.shape)\n",
    "          ind = np.moveaxis(ind, 0, -1)\n",
    "          hor_ind, ver_ind, diag_right_ind, diag_left_ind = get_lines(ind, loc)\n",
    "          # loop over each possibility\n",
    "        \n",
    "          pieces = [hor, ver, diag_right, diag_left]\n",
    "          indices = [hor_ind, ver_ind, diag_right_ind, diag_left_ind]\n",
    "        \n",
    "          #winning_loc = np.full(self.state.shape, False, dtype=bool)\n",
    "        \n",
    "          for line, index in zip(pieces, indices):\n",
    "               starts, ends, runs = get_runs(line, self.player)\n",
    "\n",
    "               # get the start and end location\n",
    "               winning = (runs >= self.N)\n",
    "               print(winning)\n",
    "               if not np.any(winning):\n",
    "                    continue\n",
    "            \n",
    "               starts_ind = starts[winning][0]\n",
    "               ends_ind = ends[winning][0]\n",
    "               indices = index[starts_ind:ends_ind]\n",
    "               #winning_loc[indices[:,0], indices[:,1]] = True\n",
    "               return indices\n",
    "            \n",
    "          return []\n",
    "    \n",
    "    \n",
    "     def move(self, loc):\n",
    "          i,j=loc\n",
    "          success = False\n",
    "          if self.w>i>=0 and self.h>j>=0:\n",
    "               if self.state[i,j]==0:\n",
    "\n",
    "                    # make a move\n",
    "                    self.state[i,j]=self.player\n",
    "\n",
    "                    # if pie rule is enabled\n",
    "                    if self.pie_rule:\n",
    "                         if self.n_moves==1:\n",
    "                              self.state[tuple(self.last_move)]=-self.player\n",
    "                              self.switched_side=False\n",
    "                    \n",
    "                         elif self.n_moves==0:\n",
    "                              # pie rule, make first move 0.5\n",
    "                              # this is to let the neural net know\n",
    "                              self.state[i,j]=self.player/2.0\n",
    "                              self.switched_side=False\n",
    "                         \n",
    "                    success = True\n",
    "\n",
    "               # switching side\n",
    "               elif self.pie_rule and self.state[i,j] == -self.player/2.0:\n",
    "\n",
    "                    # make a move\n",
    "                    self.state[i,j]=self.player\n",
    "                    self.switched_side=True\n",
    "\n",
    "                    success = True\n",
    "\n",
    "                         \n",
    "               \n",
    "\n",
    "          if success:\n",
    "               self.n_moves += 1\n",
    "               self.last_move = tuple((i,j))\n",
    "               self.score = self.get_score()\n",
    "\n",
    "               # if game is not over, switch player\n",
    "               if self.score is None:\n",
    "                    self.player *= -1\n",
    "               \n",
    "               return True\n",
    "\n",
    "          return False\n",
    "    \n",
    "    \n",
    "     def available_moves(self):\n",
    "          indices = np.moveaxis(np.indices(self.state.shape), 0, -1)\n",
    "          return indices[np.abs(self.state) != 1]\n",
    "\n",
    "     def available_mask(self):\n",
    "          return (np.abs(self.state) != 1).astype(np.uint8)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": false,
    "editable": false,
    "run_control": {
     "frozen": true
    }
   },
   "source": [
    "### The Game Play Class"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "code_folding": [
     0
    ],
    "deletable": false,
    "editable": false
   },
   "outputs": [],
   "source": [
    "class Play:\n",
    "    \n",
    "    def __init__(self, game, player1=None, player2=None, name='game'):\n",
    "        self.original_game=game\n",
    "        self.game=copy(game)\n",
    "        self.player1=player1\n",
    "        self.player2=player2\n",
    "        self.player=self.game.player\n",
    "        self.end=False\n",
    "        self.play()\n",
    "\n",
    "    def reset(self):\n",
    "        self.game=copy(self.original_game)\n",
    "        self.click_cid=None\n",
    "        self.end=False\n",
    "        \n",
    "    def play(self, name='Game'):\n",
    "        \n",
    "        self.reset()\n",
    "        \n",
    "        if self.game.w * self.game.h <25:\n",
    "            figsize=(self.game.w/1.6, self.game.h/1.6)\n",
    "        else:\n",
    "            figsize=(self.game.w/2.1, self.game.h/2.1)\n",
    "        \n",
    "        self.fig=plt.figure(name, figsize=figsize)\n",
    "        if self.game.w * self.game.h <25:\n",
    "            self.fig.subplots_adjust(.2,.2,1,1)\n",
    "        else:\n",
    "            self.fig.subplots_adjust(.1,.1,1,1)\n",
    "            \n",
    "        self.fig.show()\n",
    "        w,h=self.game.size\n",
    "        self.ax=self.fig.gca()\n",
    "        self.ax.grid()\n",
    "        # remove hovering coordinate tooltips\n",
    "        self.ax.format_coord = lambda x, y: ''\n",
    "        self.ax.set_xlim([-.5,w-.5])\n",
    "        self.ax.set_ylim([-.5,h-.5])\n",
    "        self.ax.set_xticks(np.arange(0, w, 1))\n",
    "        self.ax.set_yticks(np.arange(0, h, 1))\n",
    "        self.ax.set_aspect('equal')\n",
    "    \n",
    "        for loc in ['top', 'right', 'bottom', 'left']:\n",
    "            self.ax.spines[loc].set_visible(False)\n",
    "\n",
    "\n",
    "        # fully AI game\n",
    "        if self.player1 is not None and self.player2 is not None:\n",
    "            self.anim = animation.FuncAnimation(self.fig, self.draw_move, frames=self.move_generator, interval=500, repeat=False)\n",
    "            return\n",
    "        \n",
    "        # at least one human\n",
    "        if self.player1 is not None:\n",
    "            # first move from AI first\n",
    "            succeed = False\n",
    "            while not succeed:\n",
    "                loc = self.player1(self.game)\n",
    "                succeed = self.game.move(loc)\n",
    "\n",
    "            self.draw_move(loc)\n",
    "            \n",
    "        self.click_cid=self.fig.canvas.mpl_connect('button_press_event', self.click)\n",
    "\n",
    "            \n",
    "    def move_generator(self):\n",
    "        score = None\n",
    "        # game not concluded yet\n",
    "        while score is None:\n",
    "            self.player = self.game.player\n",
    "            if self.game.player == 1:\n",
    "                loc = self.player1(self.game)\n",
    "            else:\n",
    "                loc = self.player2(self.game)\n",
    "                \n",
    "            success = self.game.move(loc)\n",
    "\n",
    "            # see if game is done\n",
    "            if success:\n",
    "                score=self.game.score\n",
    "                yield loc\n",
    "                \n",
    "        \n",
    "    def draw_move(self, move=None):\n",
    "        if self.end:\n",
    "            return\n",
    "        \n",
    "        i,j=self.game.last_move if move is None else move\n",
    "        c='salmon' if self.player==1 else 'lightskyblue'\n",
    "        self.ax.scatter(i,j,s=500,marker='o',zorder=3, c=c)\n",
    "        score = self.game.score\n",
    "        self.draw_winner(score)\n",
    "        self.fig.canvas.draw()\n",
    "\n",
    "\n",
    "    def draw_winner(self, score):\n",
    "        if score is None:\n",
    "            return\n",
    "        \n",
    "        if score == -1 or score == 1:\n",
    "            locs = self.game.get_winning_loc()\n",
    "            c='darkred' if score==1 else 'darkblue'\n",
    "            self.ax.scatter(locs[:,0],locs[:,1], s=300, marker='*',c=c,zorder=4)\n",
    "\n",
    "        # try to disconnect if game is over\n",
    "        if hasattr(self, 'click_cid'):\n",
    "            self.fig.canvas.mpl_disconnect(self.click_cid)\n",
    "\n",
    "        self.end=True\n",
    "        \n",
    "    \n",
    "    def click(self,event):\n",
    "        \n",
    "        loc=(int(round(event.xdata)), int(round(event.ydata)))\n",
    "        self.player = self.game.player\n",
    "        succeed=self.game.move(loc)\n",
    "\n",
    "        if succeed:\n",
    "            self.draw_move()\n",
    "\n",
    "        else:\n",
    "            return\n",
    "        \n",
    "        if self.player1 is not None or self.player2 is not None:\n",
    "\n",
    "            succeed = False\n",
    "            self.player = self.game.player\n",
    "            while not succeed:\n",
    "                if self.game.player == 1:\n",
    "                    loc = self.player1(self.game)\n",
    "                else:\n",
    "                    loc = self.player2(self.game)\n",
    "                succeed = self.game.move(loc)\n",
    "               \n",
    "            self.draw_move()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": false,
    "editable": false,
    "run_control": {
     "frozen": true
    }
   },
   "source": [
    "### Setup a game"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "deletable": false,
    "editable": false
   },
   "outputs": [],
   "source": [
    "game_setting = {'size':(3,3), 'N':3}\n",
    "\n",
    "game = ConnectN(**game_setting)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": false,
    "editable": false,
    "run_control": {
     "frozen": true
    }
   },
   "source": [
    "### Display the Game State, Current Player and Score after a Move"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "deletable": false,
    "editable": false
   },
   "outputs": [],
   "source": [
    "game.move((0,1))\n",
    "print(game.state)\n",
    "print(game.player)\n",
    "print(game.score)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": false,
    "editable": false,
    "run_control": {
     "frozen": true
    }
   },
   "source": [
    "### Display the Game State after a Sequence of Moves"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "deletable": false,
    "editable": false
   },
   "outputs": [],
   "source": [
    "game.move((0,0)) # player -1 move\n",
    "game.move((1,1)) # player +1 move\n",
    "game.move((1,0)) # player -1 move\n",
    "game.move((2,1)) # player +1 move\n",
    "\n",
    "print(game.state)\n",
    "print(game.player)\n",
    "print(game.score)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": false,
    "editable": false,
    "run_control": {
     "frozen": true
    }
   },
   "source": [
    "### Play a Game Interactively"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "deletable": false,
    "editable": false
   },
   "outputs": [],
   "source": [
    "%matplotlib notebook\n",
    "\n",
    "gameplay=Play(ConnectN(**game_setting),  player1=None, player2=None)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": false,
    "editable": false,
    "run_control": {
     "frozen": true
    }
   },
   "source": [
    "### Initialize an Agent to Play the Game\n",
    "We need to define a policy for tic-tac-toe, that takes the game state as input, and outputs a policy and a critic\n",
    "\n",
    "## Tentative Exercise\n",
    "Code up your own policy for training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "code_folding": [
     0
    ],
    "deletable": false,
    "editable": false
   },
   "outputs": [],
   "source": [
    "class Policy(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(Policy, self).__init__()\n",
    "        # solution\n",
    "        self.conv = nn.Conv2d(1, 16, kernel_size=2, stride=1, bias=False)\n",
    "        self.size = 2*2*16\n",
    "        self.fc = nn.Linear(self.size,32)\n",
    "\n",
    "        # layers for the policy\n",
    "        self.fc_action1 = nn.Linear(32, 16)\n",
    "        self.fc_action2 = nn.Linear(16, 9)\n",
    "        \n",
    "        # layers for the critic\n",
    "        self.fc_value1 = nn.Linear(32, 8)\n",
    "        self.fc_value2 = nn.Linear(8, 1)\n",
    "        self.tanh_value = nn.Tanh()\n",
    "        \n",
    "    def forward(self, x):\n",
    "        x = x.to(device)\n",
    "        \n",
    "        # solution\n",
    "        y = F.relu(self.conv(x))\n",
    "        y = y.view(-1, self.size)\n",
    "        y = F.relu(self.fc(y))\n",
    "        \n",
    "        \n",
    "        # the action head\n",
    "        a = F.relu(self.fc_action1(y))\n",
    "        a = self.fc_action2(a)\n",
    "        # availability of moves\n",
    "        avail = (torch.abs(x.squeeze())!=1).float()\n",
    "        avail = avail.reshape(-1, 9)\n",
    "        \n",
    "        # locations where actions are not possible, we set the prob to zero\n",
    "        maxa = torch.max(a)\n",
    "        # subtract off max for numerical stability (avoids blowing up at infinity)\n",
    "        exp = avail*torch.exp(a-maxa)\n",
    "        prob = exp/torch.sum(exp)\n",
    "        \n",
    "        \n",
    "        # the value head\n",
    "        value = F.relu(self.fc_value1(y))\n",
    "        value = self.tanh_value(self.fc_value2(value))\n",
    "        return prob.view(3,3), value"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "deletable": false,
    "editable": false
   },
   "outputs": [],
   "source": [
    "policy = Policy()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": false,
    "editable": false,
    "run_control": {
     "frozen": true
    }
   },
   "source": [
    "### Define a Player that Uses MCTS and the Expert Policy + Critic to Play a Game\n",
    "Here we introduced a new parameter\n",
    "```\n",
    "T = temperature\n",
    "```\n",
    "This tells us how to choose the next move based on the **MCTS** results\n",
    "\n",
    "$$p_a = \\frac{N_a^{\\frac{1}{T}}}{\\sum_a N_a^{\\frac{1}{T}}}$$\n",
    "\n",
    "As $T\\rightarrow0$, we choose action with largest $N_a$. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "code_folding": [
     0,
     9
    ],
    "deletable": false,
    "editable": false
   },
   "outputs": [],
   "source": [
    "def Policy_Player_MCTS(game):\n",
    "    mytree = MCTSNode(copy(game))\n",
    "    for _ in range(50):\n",
    "        mytree.explore(policy)\n",
    "   \n",
    "    mytreenext, (v, nn_v, p, nn_p) = mytree.next(temperature=0.1)\n",
    "    \n",
    "    return mytreenext.game.last_move\n",
    "\n",
    "def Random_Player(game):\n",
    "    return random.choice(game.available_moves())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "deletable": false,
    "editable": false
   },
   "outputs": [],
   "source": [
    "game = ConnectN(**game_setting)\n",
    "\n",
    "print(game.state)\n",
    "\n",
    "Policy_Player_MCTS(game);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": false,
    "editable": false,
    "run_control": {
     "frozen": true
    }
   },
   "source": [
    "### Play a Game against the Policy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "deletable": false,
    "editable": false
   },
   "outputs": [],
   "source": [
    "%matplotlib notebook\n",
    "\n",
    "gameplay=Play(ConnectN(**game_setting), player1=None, player2=Policy_Player_MCTS)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": false,
    "editable": false,
    "run_control": {
     "frozen": true
    }
   },
   "source": [
    "## Train the Agent\n",
    "\n",
    "Initialize our **AlphaZero** agent and optimizer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "deletable": false,
    "editable": false
   },
   "outputs": [],
   "source": [
    "game=ConnectN(**game_setting)\n",
    "policy = Policy()\n",
    "optimizer = Adam(policy.parameters(), lr=0.01, weight_decay=1e-4)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": false,
    "editable": false,
    "run_control": {
     "frozen": true
    }
   },
   "source": [
    "### Tenative exercise\n",
    "code up the alphazero loss function, defined to be\n",
    "$$L = \\sum_t \\left\\{ \\left(v^{(t)}_\\theta - z\\right)^2  - \\sum_a p^{(t)}_a \\log \\pi_\\theta(a|s_t) \\right\\} + \\textrm{constant}$$ \n",
    "I added a constant term $\\sum_t \\sum_a p^{(t)}\\log p^{(t)}$ so that when $v_\\theta^{(t)} = z$ and $p^{(t)}_a = \\pi_\\theta(a|s_t)$, $L=0$, this way we can have some metric of progress."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "code_folding": [
     5
    ],
    "deletable": false,
    "editable": false
   },
   "outputs": [],
   "source": [
    "episodes = 400\n",
    "print_every = 50\n",
    "outcomes = []\n",
    "losses = []\n",
    "\n",
    "for e in range(episodes):\n",
    "\n",
    "    mytree = MCTSNode(ConnectN(**game_setting))\n",
    "    vterm = []\n",
    "    logterm = []\n",
    "    \n",
    "    while mytree.outcome is None:\n",
    "        for _ in range(50):\n",
    "            mytree.explore(policy)\n",
    "\n",
    "        current_player = mytree.game.player\n",
    "        mytree, (v, nn_v, p, nn_p) = mytree.next()        \n",
    "        mytree.detach_mother()\n",
    "        \n",
    "        # ------------- solution --------------\n",
    "        # compute prob* log pi \n",
    "        loglist = torch.log(nn_p)*p\n",
    "        \n",
    "        # constant term to make sure if policy result = MCTS result, loss = 0\n",
    "        constant = torch.where(p>0, p*torch.log(p),torch.tensor(0.))\n",
    "        logterm.append(-torch.sum(loglist-constant))\n",
    "        \n",
    "        vterm.append(nn_v*current_player)\n",
    "        # -------------------------------------\n",
    "        \n",
    "    # we compute the \"policy_loss\" for computing gradient\n",
    "    outcome = mytree.outcome\n",
    "    outcomes.append(outcome)\n",
    "    \n",
    "    # ------------- solution --------------\n",
    "    loss = torch.sum( (torch.stack(vterm)-outcome)**2 + torch.stack(logterm) )\n",
    "    # -------------------------------------\n",
    "    \n",
    "    optimizer.zero_grad()\n",
    "    loss.backward()\n",
    "    losses.append(float(loss))\n",
    "    optimizer.step()\n",
    "    \n",
    "    print('\\r' + ' ' * 120, end='')\n",
    "    print('\\rTraining Loop:: {:04d}/{:04d}    Avg. Loss: {:3.2f}    Recent Outcomes: {}'.format(e + 1, episodes, np.mean(losses[-20:]), outcomes[-10:]), end='')\n",
    "    if (e + 1) % print_every == 0:\n",
    "        print('\\r' + ' ' * 120, end='')\n",
    "        print('\\rTraining Loop:: {:04d}/{:04d}    Avg. Loss: {:3.2f}    Recent Outcomes: {}'.format(e + 1, episodes, np.mean(losses[-20:]), outcomes[-10:]))\n",
    "        \n",
    "    del loss"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": false,
    "editable": false,
    "run_control": {
     "frozen": true
    }
   },
   "source": [
    "### Plot the Loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "deletable": false,
    "editable": false
   },
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "plt.plot(losses)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": false,
    "editable": false,
    "run_control": {
     "frozen": true
    }
   },
   "source": [
    "### Play a Game against the Trained AlphaZero Agent\n",
    "\n",
    "#### As First Player"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "deletable": false,
    "editable": false
   },
   "outputs": [],
   "source": [
    "%matplotlib notebook\n",
    "\n",
    "gameplay=Play(ConnectN(**game_setting), player1=None, player2=Policy_Player_MCTS)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": false,
    "editable": false,
    "run_control": {
     "frozen": true
    }
   },
   "source": [
    "#### As Second Player"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "deletable": false,
    "editable": false
   },
   "outputs": [],
   "source": [
    "%matplotlib notebook\n",
    "\n",
    "gameplay=Play(ConnectN(**game_setting), player2=None, player1=Policy_Player_MCTS)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": false,
    "editable": false,
    "run_control": {
     "frozen": true
    }
   },
   "source": [
    "---\n",
    "\n",
    "# Advanced TicTacToe with AlphaZero\n",
    "\n",
    "We'll now extend the previous idea to an advanced version of `TicTacToe` using AlphaZero.\n",
    "\n",
    "For code walkthrough, you can watch this [video](https://www.youtube.com/watch?v=MOIk_BbCjRw)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "MODEL_FILE = './models/mcts-tictactoe-6-6-4.pt'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": false,
    "editable": false,
    "run_control": {
     "frozen": true
    }
   },
   "source": [
    "### Initialize a game"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "deletable": false,
    "editable": false
   },
   "outputs": [],
   "source": [
    "game_setting = {'size':(6,6), 'N':4, 'pie_rule':True}\n",
    "\n",
    "game = ConnectN(**game_setting)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "deletable": false,
    "editable": false
   },
   "outputs": [],
   "source": [
    "%matplotlib notebook\n",
    "\n",
    "gameplay=Play(ConnectN(**game_setting), player1=None, player2=None)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": false,
    "editable": false,
    "run_control": {
     "frozen": true
    }
   },
   "source": [
    "### Define the Policy\n",
    "\n",
    "See if you can train it under `1000` games and with only `1000` steps of exploration in each move."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "code_folding": [
     0
    ],
    "deletable": false,
    "editable": false
   },
   "outputs": [],
   "source": [
    "class Policy(nn.Module):\n",
    "\n",
    "    def __init__(self, game):\n",
    "        super(Policy, self).__init__()\n",
    "\n",
    "        # input = 6x6 board\n",
    "        # convert to 5x5x8\n",
    "        self.conv1 = nn.Conv2d(1, 16, kernel_size=2, stride=1, bias=False)\n",
    "        # 5x5x16 to 3x3x32\n",
    "        self.conv2 = nn.Conv2d(16, 32, kernel_size=3, stride=1, bias=False)\n",
    "\n",
    "        self.size=3*3*32\n",
    "        \n",
    "        # the part for actions\n",
    "        self.fc_action1 = nn.Linear(self.size, self.size//4)\n",
    "        self.fc_action2 = nn.Linear(self.size//4, 36)\n",
    "        \n",
    "        # the part for the value function\n",
    "        self.fc_value1 = nn.Linear(self.size, self.size//6)\n",
    "        self.fc_value2 = nn.Linear(self.size//6, 1)\n",
    "        self.tanh_value = nn.Tanh()\n",
    "        \n",
    "    def forward(self, x):\n",
    "        x = x.to(device)\n",
    "        \n",
    "        y = F.leaky_relu(self.conv1(x))\n",
    "        y = F.leaky_relu(self.conv2(y))\n",
    "        y = y.view(-1, self.size)\n",
    "        \n",
    "        # action head\n",
    "        a = self.fc_action2(F.leaky_relu(self.fc_action1(y)))\n",
    "        \n",
    "        avail = (torch.abs(x.squeeze())!=1).float()\n",
    "        avail = avail.reshape(-1, 36)\n",
    "        maxa = torch.max(a)\n",
    "        exp = avail*torch.exp(a-maxa)\n",
    "        prob = exp/torch.sum(exp)\n",
    "        \n",
    "        # value head\n",
    "        value = self.tanh_value(self.fc_value2(F.leaky_relu( self.fc_value1(y) )))\n",
    "        return prob.view(6,6), value"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "deletable": false,
    "editable": false
   },
   "outputs": [],
   "source": [
    "game_setting = {'size':(6,6), 'N':4}\n",
    "\n",
    "game = ConnectN(**game_setting)\n",
    "\n",
    "policy = Policy(game)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": false,
    "editable": false,
    "run_control": {
     "frozen": true
    }
   },
   "source": [
    "### Define a MCTS Player for Play"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "code_folding": [
     0,
     9
    ],
    "deletable": false,
    "editable": false
   },
   "outputs": [],
   "source": [
    "def Policy_Player_MCTS(game):\n",
    "    mytree = MCTSNode(copy(game))\n",
    "    for _ in range(1000):\n",
    "        mytree.explore(policy)\n",
    "       \n",
    "    mytreenext, (v, nn_v, p, nn_p) = mytree.next(temperature=0.1)\n",
    "    \n",
    "    return mytreenext.game.last_move\n",
    "\n",
    "def Random_Player(game):\n",
    "    return random.choice(game.available_moves())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": false,
    "editable": false,
    "run_control": {
     "frozen": true
    }
   },
   "source": [
    "### Play a Game against a Random Policy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "deletable": false,
    "editable": false
   },
   "outputs": [],
   "source": [
    "%matplotlib notebook\n",
    "\n",
    "gameplay=Play(ConnectN(**game_setting), player1=Policy_Player_MCTS, player2=None)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": false,
    "editable": false,
    "run_control": {
     "frozen": true
    }
   },
   "source": [
    "### Training\n",
    "\n",
    "Note: training is **VERY VERY** slow!!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "deletable": false,
    "editable": false
   },
   "outputs": [],
   "source": [
    "game=ConnectN(**game_setting)\n",
    "policy = Policy(game)\n",
    "optimizer = Adam(policy.parameters(), lr=0.01, weight_decay=1e-5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "code_folding": [
     7
    ],
    "deletable": false,
    "editable": false
   },
   "outputs": [],
   "source": [
    "episodes = 2000\n",
    "print_every= 10\n",
    "outcomes = []\n",
    "policy_loss = []\n",
    "Nmax = 1000\n",
    "save_every = 500\n",
    "\n",
    "for e in range(episodes):\n",
    "\n",
    "    mytree = MCTSNode(game)\n",
    "    logterm = []\n",
    "    vterm = []\n",
    "    \n",
    "    while mytree.outcome is None:\n",
    "        for _ in range(Nmax):\n",
    "            mytree.explore(policy)\n",
    "            if mytree.N >= Nmax:\n",
    "                break\n",
    "            \n",
    "        current_player = mytree.game.player\n",
    "        mytree, (v, nn_v, p, nn_p) = mytree.next()\n",
    "        mytree.detach_mother()\n",
    "        \n",
    "        loglist = torch.log(nn_p)*p\n",
    "        constant = torch.where(p>0, p*torch.log(p),torch.tensor(0.))\n",
    "        logterm.append(-torch.sum(loglist-constant))\n",
    "\n",
    "        vterm.append(nn_v*current_player)\n",
    "        \n",
    "    # we compute the \"policy_loss\" for computing gradient\n",
    "    outcome = mytree.outcome\n",
    "    outcomes.append(outcome)\n",
    "    \n",
    "    loss = torch.sum( (torch.stack(vterm)-outcome)**2 + torch.stack(logterm) )\n",
    "    optimizer.zero_grad()\n",
    "    loss.backward()\n",
    "    policy_loss.append(float(loss))\n",
    "    optimizer.step()\n",
    "    \n",
    "    print('\\r' + ' ' * 120, end='')\n",
    "    print('\\rTraining Loop:: {:04d}/{:04d}    Avg. Loss: {:3.2f}    Recent Outcomes: {}'.format(e + 1, episodes, np.mean(policy_loss[-20:]), outcomes[-10:]), end='')\n",
    "    if (e + 1) % print_every == 0:\n",
    "        print('\\r' + ' ' * 120, end='')\n",
    "        print('\\rTraining Loop:: {:04d}/{:04d}    Avg. Loss: {:3.2f}    Recent Outcomes: {}'.format(e + 1, episodes, np.mean(policy_loss[-20:]), outcomes[-10:]))\n",
    "    if (e + 1) % save_every == 0:\n",
    "        torch.save(policy, MODEL_FILE)\n",
    "    \n",
    "    del loss"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": false,
    "editable": false,
    "run_control": {
     "frozen": true
    }
   },
   "source": [
    "### Watch the Trained Agent Play against the Random Agent"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "code_folding": [
     4
    ],
    "deletable": false,
    "editable": false
   },
   "outputs": [],
   "source": [
    "%matplotlib notebook\n",
    "\n",
    "challenge_policy = torch.load(MODEL_FILE)\n",
    "\n",
    "def Challenge_Player_MCTS(game):\n",
    "    mytree = MCTSNode(copy(game))\n",
    "    for _ in range(1000):\n",
    "        mytree.explore(challenge_policy)\n",
    "    mytreenext, (v, nn_v, p, nn_p) = mytree.next(temperature=0.1)\n",
    "    return mytreenext.game.last_move\n",
    "\n",
    "gameplay=Play(ConnectN(**game_setting), player2=Policy_Player_MCTS, player1=Challenge_Player_MCTS)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": false,
    "editable": false,
    "run_control": {
     "frozen": true
    }
   },
   "source": [
    "---\n",
    "\n",
    "Next: [Multi Player Tennis](./Tennis.ipynb)"
   ]
  }
 ],
 "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.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}