connect4.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# MCTS in MCTX\n",
    "\n",
    "In this example, we will use the `mctx` library to play the game of Connect 4.\n",
    "We will implement the Monte Carlo Tree Search (MCTS) algorithm with random rollouts."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Game mechanics"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's start by defining some type aliases to make the code more readable."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import chex\n",
    "\n",
    "# 6x7 board\n",
    "# We use the following coordinate system:\n",
    "#   ^\n",
    "#   |\n",
    "#   |\n",
    "#   |\n",
    "# 0 +-------->\n",
    "#   0\n",
    "Board = chex.Array\n",
    "\n",
    "# Index of the column to play\n",
    "Action = chex.Array\n",
    "\n",
    "# Let's assume the game is played by players X and O.\n",
    "# 1 if player X, -1 if player O\n",
    "Player = chex.Array\n",
    "\n",
    "# Reward for the player who played the action.\n",
    "# 1 for winning, 0 for draw, -1 for losing\n",
    "Reward = chex.Array\n",
    "\n",
    "# True/False if the game is over\n",
    "Done = chex.Array"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, we create a class defining the game state at a given time."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "@chex.dataclass\n",
    "class Env:\n",
    "    board: Board\n",
    "    player: Player\n",
    "    done: Done\n",
    "    reward: Reward"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To visualize the game state, we define a function that prints the board."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "from rich import print\n",
    "\n",
    "BOARD_STRING = \"\"\"\n",
    " ? | ? | ? | ? | ? | ? | ?\n",
    "---|---|---|---|---|---|---\n",
    " ? | ? | ? | ? | ? | ? | ?\n",
    "---|---|---|---|---|---|---\n",
    " ? | ? | ? | ? | ? | ? | ?\n",
    "---|---|---|---|---|---|---\n",
    " ? | ? | ? | ? | ? | ? | ?\n",
    "---|---|---|---|---|---|---\n",
    " ? | ? | ? | ? | ? | ? | ?\n",
    "---|---|---|---|---|---|---\n",
    " ? | ? | ? | ? | ? | ? | ?\n",
    "---|---|---|---|---|---|---\n",
    " 1   2   3   4   5   6   7\n",
    "\"\"\"\n",
    "\n",
    "def print_board(board: Board):\n",
    "    board_str = BOARD_STRING\n",
    "    for i in reversed(range(board.shape[0])):\n",
    "        for j in range(board.shape[1]):\n",
    "            board_str = board_str.replace('?', '[green]X[/green]' if board[i, j] == 1 else '[red]O[/red]' if board[i, j] == -1 else ' ', 1)\n",
    "    print(board_str)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's write a function that checks if the game is over.\n",
    "We iterate over all horizontal/vertical/diagonal/anti-diagonal lines and check if there are 4 consecutive pieces of the same color."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "import jax.numpy as jnp\n",
    "\n",
    "def horizontals(board: Board) -> chex.Array:\n",
    "    return jnp.stack([\n",
    "        board[i, j:j+4]\n",
    "        for i in range(board.shape[0])\n",
    "        for j in range(board.shape[1] - 3)\n",
    "    ])\n",
    "\n",
    "def verticals(board: Board) -> chex.Array:\n",
    "    return jnp.stack([\n",
    "        board[i:i+4, j]\n",
    "        for i in range(board.shape[0] - 3)\n",
    "        for j in range(board.shape[1])\n",
    "    ])\n",
    "\n",
    "def diagonals(board: Board) -> chex.Array:\n",
    "    return jnp.stack([\n",
    "        jnp.diag(board[i:i+4, j:j+4])\n",
    "        for i in range(board.shape[0] - 3)\n",
    "        for j in range(board.shape[1] - 3)\n",
    "    ])\n",
    "\n",
    "def antidiagonals(board: Board) -> chex.Array:\n",
    "    return jnp.stack([\n",
    "        jnp.diag(board[i:i+4, j:j+4][::-1])\n",
    "        for i in range(board.shape[0] - 3)\n",
    "        for j in range(board.shape[1] - 3)\n",
    "    ])\n",
    "\n",
    "def get_winner(board: Board) -> Player:\n",
    "    all_lines = jnp.concatenate((\n",
    "        horizontals(board),\n",
    "        verticals(board),\n",
    "        diagonals(board),\n",
    "        antidiagonals(board),\n",
    "    ))\n",
    "    # x_won and o_won are 1 if the player won, 0 otherwise\n",
    "    x_won = jnp.any(jnp.all(all_lines == 1, axis=1)).astype(jnp.int8)\n",
    "    o_won = jnp.any(jnp.all(all_lines == -1, axis=1)).astype(jnp.int8)\n",
    "    # We consider the following cases:\n",
    "    # - !x_won and !o_won -> 0 - 0 = 0 -> draw OR not finished\n",
    "    # - x_won and !o_won -> 1 - 0 = 1 -> Player 1 (X) won\n",
    "    # - !x_won and o_won -> 0 - 1 = -1 -> Player -1 (O) won\n",
    "    # - x_won and o_won -> impossible, the game would have ended earlier\n",
    "    return x_won - o_won"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, we can implement the environment dynamics:\n",
    "\n",
    "* `env_reset` creates a brand new game\n",
    "* `env_step` plays a move and returns the new state, the reward **for the player who played the move** and whether the game has ended"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "import jax.numpy as jnp\n",
    "\n",
    "def env_reset(_):\n",
    "    return Env(\n",
    "        board=jnp.zeros((6, 7), dtype=jnp.int8),\n",
    "        player=jnp.int8(1),\n",
    "        done=jnp.bool_(False),\n",
    "        reward=jnp.int8(0),\n",
    "    )"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You might be wondering why `env_reset` takes an unused argument.\n",
    "This allows us to parallelize the environment with the `jax.vmap` function:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array([1, 1, 1, 1, 1, 1, 1, 1, 1, 1], dtype=int8)"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import jax\n",
    "\n",
    "jax.vmap(env_reset)(jnp.arange(10)).player\n",
    "# 10 games have been created:"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You need to perfectly understand the `env_step` function to implement the MCTS algorithm.\n",
    "\n",
    "In particular, you need to understand that the reward returned by `env_step` is the reward **for the player who played the move**.\n",
    "This means that the reward should always be either 0 for a move that does not end the game or causes a draw, or 1 for a move that ends the game with a win for the player who played the move.\n",
    "\n",
    "A reward of -1 is reserved for punishing illegal moves.\n",
    "Since `mctx` does not support masking illegal moves below the root node, we need another way to stop the MCTS algorithm from exploring illegal moves.\n",
    "Later, we will also write a custom policy function that will never select illegal moves."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "def env_step(env: Env, action: Action) -> tuple[Env, Reward, Done]:\n",
    "    col = action\n",
    "\n",
    "    # Find the first empty row in the column.\n",
    "    # If the column is full, this will be the top row.\n",
    "    row = jnp.argmax(env.board[:, col] == 0)\n",
    "\n",
    "    # If the column is full, the move is invalid.\n",
    "    invalid_move = env.board[row, col] != 0\n",
    "\n",
    "    # Place the player's piece in the board only if the move is valid and the game is not over.\n",
    "    board = env.board.at[row, col].set(jnp.where(env.done | invalid_move, env.board[row, col], env.player))\n",
    "\n",
    "    # The reward is computed as follows:\n",
    "    # * 0 if the game is **already** over. This is to ignore nodes below terminal nodes.\n",
    "    # * -1 if the move is invalid\n",
    "    # * 1 if the move won the game for the current player\n",
    "    # * 0 if the move caused a draw\n",
    "    # * (impossible for Connect 4) -1 if the move lost the game for the current player\n",
    "    reward = jnp.where(env.done, 0, jnp.where(invalid_move, -1, get_winner(board) * env.player)).astype(jnp.int8)\n",
    "\n",
    "    # We end the game if:\n",
    "    # * the game was already over\n",
    "    # * the move won or lost the game\n",
    "    # * the move was invalid\n",
    "    # * the board is full (draw)\n",
    "    done = env.done | reward != 0 | invalid_move | jnp.all(board[-1] != 0)\n",
    "\n",
    "    env = Env(\n",
    "        board=board,\n",
    "        # switch player\n",
    "        player=jnp.where(done, env.player, -env.player),\n",
    "        done=done,\n",
    "        reward=reward,\n",
    "    )\n",
    "\n",
    "    return env, reward, done"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## The MCTS algorithm\n",
    "\n",
    "We can now implement the MCTS algorithm."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's start with some helper functions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "def valid_action_mask(env: Env) -> chex.Array:\n",
    "    '''\n",
    "    Computes which actions are valid in the current state.\n",
    "    Returns an array of booleans, indicating which columns are not full.\n",
    "    In case the game is over, all columns are considered invalid.\n",
    "    '''\n",
    "    return jnp.where(env.done, jnp.array([False] * env.board.shape[1]), env.board[-1] == 0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "def winning_action_mask(env: Env, player: Player) -> chex.Array:\n",
    "    '''\n",
    "    Finds all actions that would immediately win the game for the given player.\n",
    "    '''\n",
    "    # Override the next player in the environment with the given player.\n",
    "    env = Env(board=env.board, player=player, done=env.done, reward=env.reward)\n",
    "\n",
    "    # Play all actions and check the reward.\n",
    "    # Remember that the reward is for the current player, so we expect it to be 1.\n",
    "    env, reward, done = jax.vmap(env_step, (None, 0))(env, jnp.arange(7, dtype=jnp.int8))\n",
    "    return reward == 1"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, we define a policy function. The policy function is used in two places:\n",
    "\n",
    "* during random rollouts to select the next action to play\n",
    "* during node expansion as the prior distribution\n",
    "\n",
    "We could have used a uniform distribution for both, but we can do better.\n",
    "Let's use three simple heuristics:\n",
    "\n",
    "* never select illegal moves\n",
    "* always select winning moves if available\n",
    "* always block opponent's winning moves if we have no winning move\n",
    "\n",
    "Our policy function implementation returns unnormalized logits,\n",
    "so we cannot return pure 0% and 100% probabilities (using `inf` and `-inf` causes downstream numerical issues).\n",
    "Instead, we add `100` to the logits of a given action which we want to prioritize:\n",
    "\n",
    "* the lowest priority of `0` is assigned to illegal moves\n",
    "* a medium priority of `100` is assigned to legal moves\n",
    "* a higher priority of `200` is assigned to the opponent's winning moves\n",
    "* the highest priority of `300` is assigned to our winning moves"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "def policy_function(env: Env) -> chex.Array:\n",
    "    return sum((\n",
    "        valid_action_mask(env).astype(jnp.float32) * 100,\n",
    "        winning_action_mask(env, -env.player).astype(jnp.float32) * 200,\n",
    "        winning_action_mask(env, env.player).astype(jnp.float32) * 300,\n",
    "    ))"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This is how these logits translate to probabilities:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array([0.14285715, 0.14285715, 0.14285715, 0.14285715, 0.14285715,\n",
       "       0.14285715, 0.14285715], dtype=float32)"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "jax.nn.softmax(jnp.array([0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array([1., 0., 0., 0., 0., 0., 0.], dtype=float32)"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "jax.nn.softmax(jnp.array([100.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array([0., 1., 0., 0., 0., 0., 0.], dtype=float32)"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "jax.nn.softmax(jnp.array([100.0, 200.0, 0.0, 0.0, 0.0, 0.0, 0.0]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array([0., 0., 1., 0., 0., 0., 0.], dtype=float32)"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "jax.nn.softmax(jnp.array([100.0, 200.0, 300.0, 0.0, 0.0, 0.0, 0.0]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array([0. , 0. , 0.5, 0.5, 0. , 0. , 0. ], dtype=float32)"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "jax.nn.softmax(jnp.array([100.0, 200.0, 300.0, 300.0, 0.0, 0.0, 0.0]))"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You can see that a difference of 100 is enough to make a move almost certain to be selected."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The MCTS algorithm uses random rollouts to estimate the value of each state.\n",
    "We can implement this with a simple function that plays moves according to our policy function until the game ends."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "def rollout(env: Env, rng_key: chex.PRNGKey) -> Reward:\n",
    "    '''\n",
    "    Plays a game until the end and returns the reward from the perspective of the initial player.\n",
    "    '''\n",
    "    def cond(a):\n",
    "        env, key = a\n",
    "        return ~env.done\n",
    "    def step(a):\n",
    "        env, key = a\n",
    "        key, subkey = jax.random.split(key)\n",
    "        action = jax.random.categorical(subkey, policy_function(env))\n",
    "        env, reward, done = env_step(env, action)\n",
    "        return env, key\n",
    "    leaf, key = jax.lax.while_loop(cond, step, (env, rng_key))\n",
    "    # The leaf reward is from the perspective of the last player.\n",
    "    # We negate it if the last player is not the initial player.\n",
    "    return leaf.reward * leaf.player * env.player"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The value function is simply the result of the rollout."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "def value_function(env: Env, rng_key: chex.PRNGKey) -> chex.Array:\n",
    "    return rollout(env, rng_key).astype(jnp.float32)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Running mctx\n",
    "\n",
    "Finally, we get to use the `mctx` library.\n",
    "We need two function:\n",
    "\n",
    "* `root_fn` returns the root node of the MCTS tree\n",
    "* `recurrent_fn` expands a new node given a parent node and an action"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "import mctx\n",
    "\n",
    "def root_fn(env: Env, rng_key: chex.PRNGKey) -> mctx.RootFnOutput:\n",
    "    return mctx.RootFnOutput(\n",
    "        prior_logits=policy_function(env),\n",
    "        value=value_function(env, rng_key),\n",
    "        # We will use the `embedding` field to store the environment.\n",
    "        embedding=env,\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "def recurrent_fn(params, rng_key, action, embedding):\n",
    "    # Extract the environment from the embedding.\n",
    "    env = embedding\n",
    "\n",
    "    # Play the action.\n",
    "    env, reward, done = env_step(env, action)\n",
    "\n",
    "    # Create the new MCTS node.\n",
    "    recurrent_fn_output = mctx.RecurrentFnOutput(\n",
    "        # reward for playing `action`\n",
    "        reward=reward,\n",
    "        # discount explained in the next section\n",
    "        discount=jnp.where(done, 0, -1).astype(jnp.float32),\n",
    "        # prior for the new state\n",
    "        prior_logits=policy_function(env),\n",
    "        # value for the new state\n",
    "        value=jnp.where(done, 0, value_function(env, rng_key)).astype(jnp.float32),\n",
    "    )\n",
    "\n",
    "    # Return the new node and the new environment.\n",
    "    return recurrent_fn_output, env"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The `discount` field is used to flip the sign of the reward and value at each tree level.\n",
    "It would be possible to implement this in the environment dynamics, but it is more convenient to do it here.\n",
    "By flipping the sign, MCTS selects the best action for the current player of each turn.\n",
    "\n",
    "Note that we set `discount` to `0` when the game is over.\n",
    "This discards all the rewards and values after the end of the game."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And now we can run MCTS!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "import functools\n",
    "from typing import Optional\n",
    "\n",
    "@functools.partial(jax.jit, static_argnums=(2,))\n",
    "def run_mcts(rng_key: chex.PRNGKey, env: Env, num_simulations: int, tree: Optional[mctx.Tree] = None) -> chex.Array:\n",
    "    batch_size = 1\n",
    "    key1, key2 = jax.random.split(rng_key)\n",
    "    policy_output = mctx.alphazero_policy(\n",
    "        # params can be used to pass additional data to the recurrent_fn like neural network weights\n",
    "        params=None,\n",
    "\n",
    "        rng_key=key1,\n",
    "\n",
    "        # create a batch of environments (in this case, a batch of size 1)\n",
    "        root=jax.vmap(root_fn, (None, 0))(env, jax.random.split(key2, batch_size)),\n",
    "\n",
    "        # automatically vectorize the recurrent_fn\n",
    "        recurrent_fn=jax.vmap(recurrent_fn, (None, None, 0, 0)),\n",
    "\n",
    "        num_simulations=num_simulations,\n",
    "\n",
    "        # we limit the depth of the search tree to 42, since we know that Connect Four can't last longer\n",
    "        max_depth=42,\n",
    "        max_nodes=int(num_simulations * 1.5),\n",
    "        search_tree=tree,\n",
    "\n",
    "        # our value is in the range [-1, 1], so we can use the min_max qtransform to map it to [0, 1]\n",
    "        qtransform=functools.partial(mctx.qtransform_by_min_max, min_value=-1, max_value=1),\n",
    "\n",
    "        # Dirichlet noise is used for exploration which we don't need in this example (we aren't training)\n",
    "        dirichlet_fraction=0.0,\n",
    "        temperature=0.0\n",
    "    )\n",
    "    return policy_output"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's play the middle column two times and see what MCTS does:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array([[0.016, 0.097, 0.189, 0.436, 0.177, 0.048, 0.037]], dtype=float32)"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "env = env_reset(0)\n",
    "env, reward, done = env_step(env, 3)\n",
    "env, reward, done = env_step(env, 3)\n",
    "policy_output = run_mcts(jax.random.PRNGKey(0), env, 1000)\n",
    "policy_output.action_weights"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAiMAAAGdCAYAAADAAnMpAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjguMiwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy8g+/7EAAAACXBIWXMAAA9hAAAPYQGoP6dpAAAZgklEQVR4nO3df6zVdf3A8de90L1XBa4oeRG8cvNHISkX5codmll5kxyz2PpBzuLu5twqNO3OJlTjai4vlTpMGIRlNZuDamk/VMxuQnNdh0Isf6SlySDtXmC1e/G6Xdy95/sH6zK+gnLwcl9ceDy2s8mHz+ec13nPeZ9+zud8bkmhUCgEAECS0uwBAICjmxgBAFKJEQAglRgBAFKJEQAglRgBAFKJEQAglRgBAFKNzB7gQPT398err74ao0ePjpKSkuxxAIADUCgUYufOnTFhwoQoLd3/+Y9hESOvvvpqVFdXZ48BAByErVu3ximnnLLfvx8WMTJ69OiI2P1mxowZkzwNAHAguru7o7q6euDn+P4Mixj530czY8aMESMAMMy83SUWLmAFAFKJEQAglRgBAFKJEQAglRgBAFKJEQAglRgBAFKJEQAglRgBAFKJEQAglRgBAFKJEQAglRgBAFKJEQAg1cjsAYDDU82CB7NHGHSbF8/OHgHYB2dGAIBUYgQASCVGAIBUYgQASCVGAIBUYgQASCVGAIBUYgQASCVGAIBUYgQASCVGAIBUYgQASCVGAIBUYgQASCVGAIBUYgQASCVGAIBUYgQASCVGAIBUYgQASCVGAIBUYgQASCVGAIBUYgQASHVQMbJs2bKoqamJioqKqK+vj/Xr1x/QcatWrYqSkpKYM2fOwbwsAHAEKjpGVq9eHc3NzdHS0hIbN26M2tramDVrVmzbtu0tj9u8eXPccMMNcdFFFx30sADAkafoGLnjjjvi6quvjqamppgyZUqsWLEijj322Ljnnnv2e0xfX19ceeWVcfPNN8dpp532jgYGAI4sRcXIrl27YsOGDdHQ0LDnCUpLo6GhIdrb2/d73Le+9a046aST4qqrrjqg1+nt7Y3u7u69HgDAkamoGNmxY0f09fVFVVXVXturqqqio6Njn8c8/vjj8aMf/SjuvvvuA36d1tbWqKysHHhUV1cXMyYAMIwc0m/T7Ny5Mz7/+c/H3XffHePGjTvg4xYuXBhdXV0Dj61btx7CKQGATCOL2XncuHExYsSI6Ozs3Gt7Z2dnjB8//k37v/TSS7F58+a4/PLLB7b19/fvfuGRI+OFF16I008//U3HlZeXR3l5eTGjAQDDVFFnRsrKymL69OnR1tY2sK2/vz/a2tpi5syZb9p/8uTJ8fTTT8emTZsGHh//+Mfjwx/+cGzatMnHLwBAcWdGIiKam5ujsbEx6urqYsaMGbFkyZLo6emJpqamiIiYN29eTJw4MVpbW6OioiLOPvvsvY4//vjjIyLetB0AODoVHSNz586N7du3x6JFi6KjoyOmTZsWa9asGbiodcuWLVFa6sauAMCBKSkUCoXsId5Od3d3VFZWRldXV4wZMyZ7HDgq1Cx4MHuEQbd58ezsEeCocqA/v53CAABSiREAIJUYAQBSiREAIJUYAQBSiREAIJUYAQBSiREAIJUYAQBSiREAIJUYAQBSiREAIJUYAQBSiREAIJUYAQBSiREAIJUYAQBSiREAIJUYAQBSiREAIJUYAQBSiREAIJUYAQBSiREAIJUYAQBSiREAIJUYAQBSiREAIJUYAQBSiREAIJUYAQBSiREAIJUYAQBSiREAIJUYAQBSiREAIJUYAQBSiREAIJUYAQBSiREAIJUYAQBSiREAIJUYAQBSiREAIJUYAQBSiREAIJUYAQBSiREAIJUYAQBSiREAIJUYAQBSiREAIJUYAQBSiREAIJUYAQBSiREAIJUYAQBSiREAIJUYAQBSiREAIJUYAQBSiREAIJUYAQBSiREAIJUYAQBSiREAIJUYAQBSiREAIJUYAQBSiREAINVBxciyZcuipqYmKioqor6+PtavX7/ffX/1q19FXV1dHH/88XHcccfFtGnT4t577z3ogQGAI0vRMbJ69epobm6OlpaW2LhxY9TW1sasWbNi27Zt+9z/hBNOiG984xvR3t4ef/3rX6OpqSmamprikUceecfDAwDDX0mhUCgUc0B9fX2cf/75sXTp0oiI6O/vj+rq6rj22mtjwYIFB/Qc5513XsyePTtuueWWA9q/u7s7Kisro6urK8aMGVPMuMBBqlnwYPYIg27z4tnZI8BR5UB/fhd1ZmTXrl2xYcOGaGho2PMEpaXR0NAQ7e3tb3t8oVCItra2eOGFF+KDH/zgfvfr7e2N7u7uvR4AwJGpqBjZsWNH9PX1RVVV1V7bq6qqoqOjY7/HdXV1xahRo6KsrCxmz54dd911V3z0ox/d7/6tra1RWVk58Kiuri5mTABgGBmSb9OMHj06Nm3aFE8++WR8+9vfjubm5li7du1+91+4cGF0dXUNPLZu3ToUYwIACUYWs/O4ceNixIgR0dnZudf2zs7OGD9+/H6PKy0tjTPOOCMiIqZNmxZ/+9vforW1NT70oQ/tc//y8vIoLy8vZjQAYJgq6sxIWVlZTJ8+Pdra2ga29ff3R1tbW8ycOfOAn6e/vz96e3uLeWkA4AhV1JmRiIjm5uZobGyMurq6mDFjRixZsiR6enqiqakpIiLmzZsXEydOjNbW1ojYff1HXV1dnH766dHb2xsPPfRQ3HvvvbF8+fLBfScAwLBUdIzMnTs3tm/fHosWLYqOjo6YNm1arFmzZuCi1i1btkRp6Z4TLj09PfHlL385/vWvf8UxxxwTkydPjp/97Gcxd+7cwXsXAMCwVfR9RjK4zwgMPfcZAd6pQ3KfEQCAwSZGAIBUYgQASCVGAIBUYgQASCVGAIBUYgQASCVGAIBUYgQASCVGAIBUYgQASCVGAIBUYgQASCVGAIBUYgQASCVGAIBUYgQASCVGAIBUYgQASCVGAIBUYgQASCVGAIBUYgQASCVGAIBUYgQASCVGAIBUYgQASCVGAIBUYgQASCVGAIBUYgQASCVGAIBUYgQASCVGAIBUYgQASCVGAIBUYgQASCVGAIBUYgQASCVGAIBUYgQASCVGAIBUYgQASCVGAIBUYgQASCVGAIBUYgQASCVGAIBUYgQASCVGAIBUYgQASCVGAIBUYgQASCVGAIBUYgQASCVGAIBUYgQASCVGAIBUYgQASCVGAIBUYgQASCVGAIBUYgQASCVGAIBUYgQASCVGAIBUYgQASCVGAIBUYgQASCVGAIBUBxUjy5Yti5qamqioqIj6+vpYv379fve9++6746KLLoqxY8fG2LFjo6Gh4S33BwCOLkXHyOrVq6O5uTlaWlpi48aNUVtbG7NmzYpt27btc/+1a9fGFVdcEY899li0t7dHdXV1XHrppfHKK6+84+EBgOGvpFAoFIo5oL6+Ps4///xYunRpRET09/dHdXV1XHvttbFgwYK3Pb6vry/Gjh0bS5cujXnz5h3Qa3Z3d0dlZWV0dXXFmDFjihkXOEg1Cx7MHmHQbV48O3sEOKoc6M/vos6M7Nq1KzZs2BANDQ17nqC0NBoaGqK9vf2AnuP111+PN954I0444YT97tPb2xvd3d17PQCAI1NRMbJjx47o6+uLqqqqvbZXVVVFR0fHAT3HjTfeGBMmTNgraP6/1tbWqKysHHhUV1cXMyYAMIwM6bdpFi9eHKtWrYr7778/Kioq9rvfwoULo6ura+CxdevWIZwSABhKI4vZedy4cTFixIjo7Ozca3tnZ2eMHz/+LY+97bbbYvHixfGHP/whpk6d+pb7lpeXR3l5eTGjAQDDVFFnRsrKymL69OnR1tY2sK2/vz/a2tpi5syZ+z3uu9/9btxyyy2xZs2aqKurO/hpAYAjTlFnRiIimpubo7GxMerq6mLGjBmxZMmS6OnpiaampoiImDdvXkycODFaW1sjIuI73/lOLFq0KO67776oqakZuLZk1KhRMWrUqEF8KwDAcFR0jMydOze2b98eixYtio6Ojpg2bVqsWbNm4KLWLVu2RGnpnhMuy5cvj127dsWnPvWpvZ6npaUlbrrppnc2PQAw7BV9n5EM7jMCQ899RoB36pDcZwQAYLCJEQAgVdHXjMCRzscTAEPLmREAIJUYAQBSiREAIJUYAQBSiREAIJUYAQBSiREAIJUYAQBSuekZwNtwIzw4tJwZAQBSiREAIJUYAQBSiREAIJUYAQBSiREAIJUYAQBSiREAIJUYAQBSiREAIJUYAQBSiREAIJUYAQBSiREAIJUYAQBSiREAIJUYAQBSiREAIJUYAQBSiREAIJUYAQBSiREAIJUYAQBSiREAIJUYAQBSiREAIJUYAQBSiREAIJUYAQBSiREAIJUYAQBSiREAIJUYAQBSiREAIJUYAQBSiREAIJUYAQBSiREAIJUYAQBSiREAIJUYAQBSiREAIJUYAQBSiREAIJUYAQBSiREAIJUYAQBSiREAIJUYAQBSiREAIJUYAQBSiREAIJUYAQBSiREAIJUYAQBSiREAIJUYAQBSHVSMLFu2LGpqaqKioiLq6+tj/fr1+9332WefjU9+8pNRU1MTJSUlsWTJkoOdFQA4AhUdI6tXr47m5uZoaWmJjRs3Rm1tbcyaNSu2bdu2z/1ff/31OO2002Lx4sUxfvz4dzwwAHBkKTpG7rjjjrj66qujqakppkyZEitWrIhjjz027rnnnn3uf/7558f3vve9+OxnPxvl5eXveGAA4MhSVIzs2rUrNmzYEA0NDXueoLQ0Ghoaor29fdCHAwCOfCOL2XnHjh3R19cXVVVVe22vqqqK559/ftCG6u3tjd7e3oE/d3d3D9pzAwCHl8Py2zStra1RWVk58Kiurs4eCQA4RIqKkXHjxsWIESOis7Nzr+2dnZ2DenHqwoULo6ura+CxdevWQXtuAODwUtTHNGVlZTF9+vRoa2uLOXPmREREf39/tLW1xTXXXDNoQ5WXl7vYNUHNggezRxh0mxfPzh4BgLdRVIxERDQ3N0djY2PU1dXFjBkzYsmSJdHT0xNNTU0RETFv3ryYOHFitLa2RsTui16fe+65gX9+5ZVXYtOmTTFq1Kg444wzBvGtAADDUdExMnfu3Ni+fXssWrQoOjo6Ytq0abFmzZqBi1q3bNkSpaV7Pv159dVX49xzzx3482233Ra33XZbXHzxxbF27dp3/g4AgGGt6BiJiLjmmmv2+7HM/w+MmpqaKBQKB/MyAMBR4LD8Ng0AcPQQIwBAKjECAKQSIwBAKjECAKQSIwBAKjECAKQSIwBAKjECAKQSIwBAKjECAKQSIwBAKjECAKQSIwBAKjECAKQSIwBAKjECAKQSIwBAKjECAKQSIwBAKjECAKQSIwBAKjECAKQSIwBAKjECAKQSIwBAKjECAKQSIwBAKjECAKQSIwBAKjECAKQSIwBAKjECAKQSIwBAKjECAKQSIwBAqpHZAwAwPNQseDB7hEG3efHs7BEIZ0YAgGRiBABIJUYAgFRiBABIJUYAgFRiBABIJUYAgFTuMwIARXC/lcHnzAgAkEqMAACpxAgAkEqMAACpxAgAkEqMAACpxAgAkEqMAACpxAgAkEqMAACpxAgAkEqMAACpxAgAkEqMAACpxAgAkEqMAACpRmYPkK1mwYPZIwy6zYtnZ48AAAfMmREAIJUYAQBSiREAIJUYAQBSiREAIJUYAQBSiREAIJUYAQBSHVSMLFu2LGpqaqKioiLq6+tj/fr1b7n/L37xi5g8eXJUVFTEOeecEw899NBBDQsAHHmKjpHVq1dHc3NztLS0xMaNG6O2tjZmzZoV27Zt2+f+f/7zn+OKK66Iq666Kv7yl7/EnDlzYs6cOfHMM8+84+EBgOGv6Bi544474uqrr46mpqaYMmVKrFixIo499ti455579rn/nXfeGR/72Mfia1/7Wpx11llxyy23xHnnnRdLly59x8MDAMNfUb+bZteuXbFhw4ZYuHDhwLbS0tJoaGiI9vb2fR7T3t4ezc3Ne22bNWtWPPDAA/t9nd7e3ujt7R34c1dXV0REdHd3FzPuAenvfX3QnzPbwa6TtdjNOuxmHfawFrtZh92sQ/HPWygU3nK/omJkx44d0dfXF1VVVXttr6qqiueff36fx3R0dOxz/46Ojv2+Tmtra9x8881v2l5dXV3MuEetyiXZExw+rMVu1mE367CHtdjNOux2qNdh586dUVlZud+/Pyx/a+/ChQv3OpvS398f//nPf+LEE0+MkpKSxMkOXnd3d1RXV8fWrVtjzJgx2eOksQ67WYc9rMVu1mE367DHkbAWhUIhdu7cGRMmTHjL/YqKkXHjxsWIESOis7Nzr+2dnZ0xfvz4fR4zfvz4ovaPiCgvL4/y8vK9th1//PHFjHrYGjNmzLD9l2owWYfdrMMe1mI367CbddhjuK/FW50R+Z+iLmAtKyuL6dOnR1tb28C2/v7+aGtri5kzZ+7zmJkzZ+61f0TEo48+ut/9AYCjS9Ef0zQ3N0djY2PU1dXFjBkzYsmSJdHT0xNNTU0RETFv3ryYOHFitLa2RkTEddddFxdffHHcfvvtMXv27Fi1alU89dRTsXLlysF9JwDAsFR0jMydOze2b98eixYtio6Ojpg2bVqsWbNm4CLVLVu2RGnpnhMuF1xwQdx3333xzW9+M77+9a/HmWeeGQ888ECcffbZg/cuhoHy8vJoaWl508dPRxvrsJt12MNa7GYddrMOexxNa1FSeLvv2wAAHEJ+Nw0AkEqMAACpxAgAkEqMAACpxMgQWLZsWdTU1ERFRUXU19fH+vXrs0cacn/605/i8ssvjwkTJkRJSclb/m6iI1lra2ucf/75MXr06DjppJNizpw58cILL2SPNeSWL18eU6dOHbiZ08yZM+Phhx/OHivd4sWLo6SkJK6//vrsUYbcTTfdFCUlJXs9Jk+enD1WildeeSU+97nPxYknnhjHHHNMnHPOOfHUU09lj3VIiZFDbPXq1dHc3BwtLS2xcePGqK2tjVmzZsW2bduyRxtSPT09UVtbG8uWLcseJdW6deti/vz58cQTT8Sjjz4ab7zxRlx66aXR09OTPdqQOuWUU2Lx4sWxYcOGeOqpp+IjH/lIfOITn4hnn302e7Q0Tz75ZPzgBz+IqVOnZo+S5v3vf3/8+9//Hng8/vjj2SMNuf/+979x4YUXxrve9a54+OGH47nnnovbb789xo4dmz3aoVXgkJoxY0Zh/vz5A3/u6+srTJgwodDa2po4Va6IKNx///3ZYxwWtm3bVoiIwrp167JHSTd27NjCD3/4w+wxUuzcubNw5plnFh599NHCxRdfXLjuuuuyRxpyLS0thdra2uwx0t14442FD3zgA9ljDDlnRg6hXbt2xYYNG6KhoWFgW2lpaTQ0NER7e3viZBwuurq6IiLihBNOSJ4kT19fX6xatSp6enqO2l8TMX/+/Jg9e/Ze/604Gv3jH/+ICRMmxGmnnRZXXnllbNmyJXukIfeb3/wm6urq4tOf/nScdNJJce6558bdd9+dPdYhJ0YOoR07dkRfX9/A3Wn/p6qqKjo6OpKm4nDR398f119/fVx44YVH3R2JIyKefvrpGDVqVJSXl8cXv/jFuP/++2PKlCnZYw25VatWxcaNGwd+hcbRqr6+Pn7yk5/EmjVrYvny5fHyyy/HRRddFDt37swebUj985//jOXLl8eZZ54ZjzzySHzpS1+Kr3zlK/HTn/40e7RDqujbwQODY/78+fHMM88clZ+LR0S8733vi02bNkVXV1f88pe/jMbGxli3bt1RFSRbt26N6667Lh599NGoqKjIHifVZZddNvDPU6dOjfr6+pg0aVL8/Oc/j6uuuipxsqHV398fdXV1ceutt0ZExLnnnhvPPPNMrFixIhobG5OnO3ScGTmExo0bFyNGjIjOzs69tnd2dsb48eOTpuJwcM0118Tvfve7eOyxx+KUU07JHidFWVlZnHHGGTF9+vRobW2N2trauPPOO7PHGlIbNmyIbdu2xXnnnRcjR46MkSNHxrp16+L73/9+jBw5Mvr6+rJHTHP88cfHe9/73njxxRezRxlSJ5988puC/KyzzjriP7ISI4dQWVlZTJ8+Pdra2ga29ff3R1tb21H72fjRrlAoxDXXXBP3339//PGPf4z3vOc92SMdNvr7+6O3tzd7jCF1ySWXxNNPPx2bNm0aeNTV1cWVV14ZmzZtihEjRmSPmOa1116Ll156KU4++eTsUYbUhRde+Kav+//973+PSZMmJU00NHxMc4g1NzdHY2Nj1NXVxYwZM2LJkiXR09MTTU1N2aMNqddee22v/8N5+eWXY9OmTXHCCSfEqaeemjjZ0Jo/f37cd9998etf/zpGjx49cO1QZWVlHHPMMcnTDZ2FCxfGZZddFqeeemrs3Lkz7rvvvli7dm088sgj2aMNqdGjR7/peqHjjjsuTjzxxKPuOqIbbrghLr/88pg0aVK8+uqr0dLSEiNGjIgrrrgie7Qh9dWvfjUuuOCCuPXWW+Mzn/lMrF+/PlauXBkrV67MHu3Qyv46z9HgrrvuKpx66qmFsrKywowZMwpPPPFE9khD7rHHHitExJsejY2N2aMNqX2tQUQUfvzjH2ePNqS+8IUvFCZNmlQoKysrvPvd7y5ccsklhd///vfZYx0Wjtav9s6dO7dw8sknF8rKygoTJ04szJ07t/Diiy9mj5Xit7/9beHss88ulJeXFyZPnlxYuXJl9kiHXEmhUCgkdRAAgGtGAIBcYgQASCVGAIBUYgQASCVGAIBUYgQASCVGAIBUYgQASCVGAIBUYgQASCVGAIBUYgQASPV/yIqCzOtC/tEAAAAASUVORK5CYII=",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "plt.bar(jnp.arange(7), policy_output.action_weights.mean(axis=0))\n",
    "plt.show()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "It chose the middle column as the best move, as expected."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's write a simple script to play against MCTS:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# set to False to enable human input\n",
    "player_1_ai = True\n",
    "player_2_ai = True\n",
    "player_1_use_subtree = True\n",
    "player_2_use_subtree = False\n",
    "simulations_p1 = 2000\n",
    "simulations_p2 = 2000\n",
    "key = jax.random.PRNGKey(0)\n",
    "env = env_reset(0)\n",
    "print_board(env.board)\n",
    "tree=None\n",
    "tree2=None\n",
    "while True:\n",
    "    if player_1_ai:\n",
    "        output = run_mcts(key, env, simulations_p1, tree)\n",
    "        output.action\n",
    "        \n",
    "        if player_1_use_subtree:\n",
    "            tree = mctx.get_subtree(output.search_tree, output.action)\n",
    "        action = output.action.item()\n",
    "    else:\n",
    "        action = int(input()) - 1\n",
    "\n",
    "    if player_2_ai and player_2_use_subtree:\n",
    "        if tree2:\n",
    "            tree2 = mctx.get_subtree(tree2, jnp.array([action], dtype=jnp.int32).reshape(1, -1))\n",
    "\n",
    "    env, reward, done = env_step(env, action)\n",
    "    print_board(env.board)\n",
    "    if done: break\n",
    "\n",
    "    if player_2_ai:\n",
    "        # you can give it more simulations to make it stronger\n",
    "        output = run_mcts(key, env, simulations_p2, tree2)\n",
    "        \n",
    "        if player_2_use_subtree:\n",
    "            tree2 = mctx.get_subtree(output.search_tree, output.action)\n",
    "        action = output.action.item()\n",
    "    else:\n",
    "        action = int(input()) - 1\n",
    "    \n",
    "    if player_1_ai and player_1_use_subtree:\n",
    "        if tree:\n",
    "            tree = mctx.get_subtree(tree, jnp.array([action], dtype=jnp.int32).reshape(1, -1))\n",
    "\n",
    "    env, reward, done = env_step(env, action)\n",
    "    print_board(env.board)\n",
    "    if done: break\n",
    "\n",
    "players = {\n",
    "    1: \"[green]X[/green]\",\n",
    "    -1: \"[red]O[/red]\",\n",
    "}\n",
    "print(reward)\n",
    "print(f\"Winner: {players[env.player.item()]}\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "At `10_000` simulations, the agent was able to beat all online Connect 4 bots I could find."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "mctx-classic",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.9"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}