colab/Colab_UnityEnvironment_2_Train.ipynb

{
 "nbformat": 4,
 "nbformat_minor": 0,
 "metadata": {
  "colab": {
   "name": "Colab-UnityEnvironment-2-Train.ipynb",
   "private_outputs": true,
   "provenance": [],
   "collapsed_sections": []
  },
  "kernelspec": {
   "name": "python3",
   "language": "python",
   "display_name": "Python 3"
  }
 },
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "pbVXrmEsLXDt"
   },
   "source": [
    "# ML-Agents Q-Learning with GridWorld\n",
    "<img src=\"https://github.com/Unity-Technologies/ml-agents/blob/release_20_docs/docs/images/gridworld.png?raw=true\" align=\"middle\" width=\"435\"/>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "WNKTwHU3d2-l"
   },
   "source": [
    "## Setup"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "htb-p1hSNX7D"
   },
   "source": [
    "#@title Install Rendering Dependencies { display-mode: \"form\" }\n",
    "#@markdown (You only need to run this code when using Colab's hosted runtime)\n",
    "\n",
    "import os\n",
    "from IPython.display import HTML, display\n",
    "\n",
    "def progress(value, max=100):\n",
    "    return HTML(\"\"\"\n",
    "        <progress\n",
    "            value='{value}'\n",
    "            max='{max}',\n",
    "            style='width: 100%'\n",
    "        >\n",
    "            {value}\n",
    "        </progress>\n",
    "    \"\"\".format(value=value, max=max))\n",
    "\n",
    "pro_bar = display(progress(0, 100), display_id=True)\n",
    "\n",
    "try:\n",
    "  import google.colab\n",
    "  INSTALL_XVFB = True\n",
    "except ImportError:\n",
    "  INSTALL_XVFB = 'COLAB_ALWAYS_INSTALL_XVFB' in os.environ\n",
    "\n",
    "if INSTALL_XVFB:\n",
    "  with open('frame-buffer', 'w') as writefile:\n",
    "    writefile.write(\"\"\"#taken from https://gist.github.com/jterrace/2911875\n",
    "XVFB=/usr/bin/Xvfb\n",
    "XVFBARGS=\":1 -screen 0 1024x768x24 -ac +extension GLX +render -noreset\"\n",
    "PIDFILE=./frame-buffer.pid\n",
    "case \"$1\" in\n",
    "  start)\n",
    "    echo -n \"Starting virtual X frame buffer: Xvfb\"\n",
    "    /sbin/start-stop-daemon --start --quiet --pidfile $PIDFILE --make-pidfile --background --exec $XVFB -- $XVFBARGS\n",
    "    echo \".\"\n",
    "    ;;\n",
    "  stop)\n",
    "    echo -n \"Stopping virtual X frame buffer: Xvfb\"\n",
    "    /sbin/start-stop-daemon --stop --quiet --pidfile $PIDFILE\n",
    "    rm $PIDFILE\n",
    "    echo \".\"\n",
    "    ;;\n",
    "  restart)\n",
    "    $0 stop\n",
    "    $0 start\n",
    "    ;;\n",
    "  *)\n",
    "        echo \"Usage: /etc/init.d/xvfb {start|stop|restart}\"\n",
    "        exit 1\n",
    "esac\n",
    "exit 0\n",
    "    \"\"\")\n",
    "  !sudo apt-get update\n",
    "  pro_bar.update(progress(10, 100))\n",
    "  !sudo DEBIAN_FRONTEND=noninteractive apt install -y daemon wget gdebi-core build-essential libfontenc1 libfreetype6 xorg-dev xorg\n",
    "  pro_bar.update(progress(20, 100))\n",
    "  !wget http://security.ubuntu.com/ubuntu/pool/main/libx/libxfont/libxfont1_1.5.1-1ubuntu0.16.04.4_amd64.deb 2>&1\n",
    "  pro_bar.update(progress(30, 100))\n",
    "  !wget --output-document xvfb.deb http://security.ubuntu.com/ubuntu/pool/universe/x/xorg-server/xvfb_1.18.4-0ubuntu0.12_amd64.deb 2>&1\n",
    "  pro_bar.update(progress(40, 100))\n",
    "  !sudo dpkg -i libxfont1_1.5.1-1ubuntu0.16.04.4_amd64.deb 2>&1\n",
    "  pro_bar.update(progress(50, 100))\n",
    "  !sudo dpkg -i xvfb.deb 2>&1\n",
    "  pro_bar.update(progress(70, 100))\n",
    "  !rm libxfont1_1.5.1-1ubuntu0.16.04.4_amd64.deb\n",
    "  pro_bar.update(progress(80, 100))\n",
    "  !rm xvfb.deb\n",
    "  pro_bar.update(progress(90, 100))\n",
    "  !bash frame-buffer start\n",
    "  os.environ[\"DISPLAY\"] = \":1\"\n",
    "pro_bar.update(progress(100, 100))"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Pzj7wgapAcDs"
   },
   "source": [
    "### Installing ml-agents"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "N8yfQqkbebQ5"
   },
   "source": [
    "try:\n",
    "  import mlagents\n",
    "  print(\"ml-agents already installed\")\n",
    "except ImportError:\n",
    "  !python -m pip install -q mlagents==0.30.0\n",
    "  print(\"Installed ml-agents\")"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "jz81TWAkbuFY"
   },
   "source": [
    "## Train the GridWorld Environment with Q-Learning"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "29n3dt1Zx5ty"
   },
   "source": [
    "### What is the GridWorld Environment"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "pZhVRfdoyPmv"
   },
   "source": [
    "The [GridWorld](https://github.com/Unity-Technologies/ml-agents/blob/release_20_docs/docs/Learning-Environment-Examples.md#gridworld) Environment is a simple Unity visual environment. The Agent is a blue square in a 3x3 grid that is trying to reach a green __`+`__ while avoiding a red __`x`__.\n",
    "\n",
    "The observation is an image obtained by a camera on top of the grid.\n",
    "\n",
    "The Action can be one of 5 :\n",
    " - Do not move\n",
    " - Move up\n",
    " - Move down\n",
    " - Move right\n",
    " - Move left\n",
    "\n",
    "The Agent receives a reward of _1.0_ if it reaches the green __`+`__, a penalty of _-1.0_ if it touches the red __`x`__ and a penalty of `-0.01` at every step (to force the Agent to solve the task as fast as possible)\n",
    "\n",
    "__Note__ There are 9 Agents, each in their own grid, at once in the environment. This alows for faster data collection.\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "4Gt-ZydJyJWD"
   },
   "source": [
    "### The Q-Learning Algorithm\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "KA1qOgfq0Xdv"
   },
   "source": [
    "In this Notebook, we will implement a very simple Q-Learning algorithm. We will use [pytorch](https://pytorch.org/) to do so.\n",
    "\n",
    "Below is the code to create the neural network we will use in the Notebook."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "q79rUp_Sx6A_"
   },
   "source": [
    "import torch\n",
    "from typing import Tuple\n",
    "from math import floor\n",
    "from torch.nn import Parameter\n",
    "\n",
    "\n",
    "class VisualQNetwork(torch.nn.Module):\n",
    "  def __init__(\n",
    "    self,\n",
    "    input_shape: Tuple[int, int, int],\n",
    "    encoding_size: int,\n",
    "    output_size: int\n",
    "  ):\n",
    "    \"\"\"\n",
    "    Creates a neural network that takes as input a batch of images (3\n",
    "    dimensional tensors) and outputs a batch of outputs (1 dimensional\n",
    "    tensors)\n",
    "    \"\"\"\n",
    "    super(VisualQNetwork, self).__init__()\n",
    "    height = input_shape[1]\n",
    "    width = input_shape[2]\n",
    "    initial_channels = input_shape[0]\n",
    "    conv_1_hw = self.conv_output_shape((height, width), 8, 4)\n",
    "    conv_2_hw = self.conv_output_shape(conv_1_hw, 4, 2)\n",
    "    self.final_flat = conv_2_hw[0] * conv_2_hw[1] * 32\n",
    "    self.conv1 = torch.nn.Conv2d(initial_channels, 16, [8, 8], [4, 4])\n",
    "    self.conv2 = torch.nn.Conv2d(16, 32, [4, 4], [2, 2])\n",
    "    self.dense1 = torch.nn.Linear(self.final_flat, encoding_size)\n",
    "    self.dense2 = torch.nn.Linear(encoding_size, output_size)\n",
    "\n",
    "  def forward(self, visual_obs: torch.tensor):\n",
    "    conv_1 = torch.relu(self.conv1(visual_obs))\n",
    "    conv_2 = torch.relu(self.conv2(conv_1))\n",
    "    hidden = self.dense1(conv_2.reshape([-1, self.final_flat]))\n",
    "    hidden = torch.relu(hidden)\n",
    "    hidden = self.dense2(hidden)\n",
    "    return hidden\n",
    "\n",
    "  @staticmethod\n",
    "  def conv_output_shape(\n",
    "    h_w: Tuple[int, int],\n",
    "    kernel_size: int = 1,\n",
    "    stride: int = 1,\n",
    "    pad: int = 0,\n",
    "    dilation: int = 1,\n",
    "  ):\n",
    "    \"\"\"\n",
    "    Computes the height and width of the output of a convolution layer.\n",
    "    \"\"\"\n",
    "    h = floor(\n",
    "      ((h_w[0] + (2 * pad) - (dilation * (kernel_size - 1)) - 1) / stride) + 1\n",
    "    )\n",
    "    w = floor(\n",
    "      ((h_w[1] + (2 * pad) - (dilation * (kernel_size - 1)) - 1) / stride) + 1\n",
    "    )\n",
    "    return h, w\n"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "EZoaEBAo2L0F"
   },
   "source": [
    "We will now create a few classes to help us store the data we will use to train the Q-Learning algorithm."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "L772fe2q39DO"
   },
   "source": [
    "import numpy as np\n",
    "from typing import NamedTuple, List\n",
    "\n",
    "\n",
    "class Experience(NamedTuple):\n",
    "  \"\"\"\n",
    "  An experience contains the data of one Agent transition.\n",
    "  - Observation\n",
    "  - Action\n",
    "  - Reward\n",
    "  - Done flag\n",
    "  - Next Observation\n",
    "  \"\"\"\n",
    "\n",
    "  obs: np.ndarray\n",
    "  action: np.ndarray\n",
    "  reward: float\n",
    "  done: bool\n",
    "  next_obs: np.ndarray\n",
    "\n",
    "# A Trajectory is an ordered sequence of Experiences\n",
    "Trajectory = List[Experience]\n",
    "\n",
    "# A Buffer is an unordered list of Experiences from multiple Trajectories\n",
    "Buffer = List[Experience]"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "6HsM1d5I3_Tj"
   },
   "source": [
    "Now, we can create our trainer class. The role of this trainer is to collect data from the Environment according to a Policy, and then train the Q-Network with that data."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "KkzBoRJCb18t"
   },
   "source": [
    "from mlagents_envs.environment import ActionTuple, BaseEnv\n",
    "from typing import Dict\n",
    "import random\n",
    "\n",
    "\n",
    "class Trainer:\n",
    "  @staticmethod\n",
    "  def generate_trajectories(\n",
    "    env: BaseEnv, q_net: VisualQNetwork, buffer_size: int, epsilon: float\n",
    "  ):\n",
    "    \"\"\"\n",
    "    Given a Unity Environment and a Q-Network, this method will generate a\n",
    "    buffer of Experiences obtained by running the Environment with the Policy\n",
    "    derived from the Q-Network.\n",
    "    :param BaseEnv: The UnityEnvironment used.\n",
    "    :param q_net: The Q-Network used to collect the data.\n",
    "    :param buffer_size: The minimum size of the buffer this method will return.\n",
    "    :param epsilon: Will add a random normal variable with standard deviation.\n",
    "    epsilon to the value heads of the Q-Network to encourage exploration.\n",
    "    :returns: a Tuple containing the created buffer and the average cumulative\n",
    "    the Agents obtained.\n",
    "    \"\"\"\n",
    "    # Create an empty Buffer\n",
    "    buffer: Buffer = []\n",
    "\n",
    "    # Reset the environment\n",
    "    env.reset()\n",
    "    # Read and store the Behavior Name of the Environment\n",
    "    behavior_name = list(env.behavior_specs)[0]\n",
    "    # Read and store the Behavior Specs of the Environment\n",
    "    spec = env.behavior_specs[behavior_name]\n",
    "\n",
    "    # Create a Mapping from AgentId to Trajectories. This will help us create\n",
    "    # trajectories for each Agents\n",
    "    dict_trajectories_from_agent: Dict[int, Trajectory] = {}\n",
    "    # Create a Mapping from AgentId to the last observation of the Agent\n",
    "    dict_last_obs_from_agent: Dict[int, np.ndarray] = {}\n",
    "    # Create a Mapping from AgentId to the last observation of the Agent\n",
    "    dict_last_action_from_agent: Dict[int, np.ndarray] = {}\n",
    "    # Create a Mapping from AgentId to cumulative reward (Only for reporting)\n",
    "    dict_cumulative_reward_from_agent: Dict[int, float] = {}\n",
    "    # Create a list to store the cumulative rewards obtained so far\n",
    "    cumulative_rewards: List[float] = []\n",
    "\n",
    "    while len(buffer) < buffer_size:  # While not enough data in the buffer\n",
    "      # Get the Decision Steps and Terminal Steps of the Agents\n",
    "      decision_steps, terminal_steps = env.get_steps(behavior_name)\n",
    "\n",
    "      # permute the tensor to go from NHWC to NCHW\n",
    "      order = (0, 3, 1, 2)\n",
    "      decision_steps.obs = [np.transpose(obs, order) for obs in decision_steps.obs]\n",
    "      terminal_steps.obs = [np.transpose(obs, order) for obs in terminal_steps.obs]\n",
    "\n",
    "      # For all Agents with a Terminal Step:\n",
    "      for agent_id_terminated in terminal_steps:\n",
    "        # Create its last experience (is last because the Agent terminated)\n",
    "        last_experience = Experience(\n",
    "          obs=dict_last_obs_from_agent[agent_id_terminated].copy(),\n",
    "          reward=terminal_steps[agent_id_terminated].reward,\n",
    "          done=not terminal_steps[agent_id_terminated].interrupted,\n",
    "          action=dict_last_action_from_agent[agent_id_terminated].copy(),\n",
    "          next_obs=terminal_steps[agent_id_terminated].obs[0],\n",
    "        )\n",
    "        # Clear its last observation and action (Since the trajectory is over)\n",
    "        dict_last_obs_from_agent.pop(agent_id_terminated)\n",
    "        dict_last_action_from_agent.pop(agent_id_terminated)\n",
    "        # Report the cumulative reward\n",
    "        cumulative_reward = (\n",
    "          dict_cumulative_reward_from_agent.pop(agent_id_terminated)\n",
    "          + terminal_steps[agent_id_terminated].reward\n",
    "        )\n",
    "        cumulative_rewards.append(cumulative_reward)\n",
    "        # Add the Trajectory and the last experience to the buffer\n",
    "        buffer.extend(dict_trajectories_from_agent.pop(agent_id_terminated))\n",
    "        buffer.append(last_experience)\n",
    "\n",
    "      # For all Agents with a Decision Step:\n",
    "      for agent_id_decisions in decision_steps:\n",
    "        # If the Agent does not have a Trajectory, create an empty one\n",
    "        if agent_id_decisions not in dict_trajectories_from_agent:\n",
    "          dict_trajectories_from_agent[agent_id_decisions] = []\n",
    "          dict_cumulative_reward_from_agent[agent_id_decisions] = 0\n",
    "\n",
    "        # If the Agent requesting a decision has a \"last observation\"\n",
    "        if agent_id_decisions in dict_last_obs_from_agent:\n",
    "          # Create an Experience from the last observation and the Decision Step\n",
    "          exp = Experience(\n",
    "            obs=dict_last_obs_from_agent[agent_id_decisions].copy(),\n",
    "            reward=decision_steps[agent_id_decisions].reward,\n",
    "            done=False,\n",
    "            action=dict_last_action_from_agent[agent_id_decisions].copy(),\n",
    "            next_obs=decision_steps[agent_id_decisions].obs[0],\n",
    "          )\n",
    "          # Update the Trajectory of the Agent and its cumulative reward\n",
    "          dict_trajectories_from_agent[agent_id_decisions].append(exp)\n",
    "          dict_cumulative_reward_from_agent[agent_id_decisions] += (\n",
    "            decision_steps[agent_id_decisions].reward\n",
    "          )\n",
    "        # Store the observation as the new \"last observation\"\n",
    "        dict_last_obs_from_agent[agent_id_decisions] = (\n",
    "          decision_steps[agent_id_decisions].obs[0]\n",
    "        )\n",
    "\n",
    "      # Generate an action for all the Agents that requested a decision\n",
    "      # Compute the values for each action given the observation\n",
    "      actions_values = (\n",
    "        q_net(torch.from_numpy(decision_steps.obs[0])).detach().numpy()\n",
    "      )\n",
    "      # Add some noise with epsilon to the values\n",
    "      actions_values += epsilon * (\n",
    "        np.random.randn(actions_values.shape[0], actions_values.shape[1])\n",
    "      ).astype(np.float32)\n",
    "      # Pick the best action using argmax\n",
    "      actions = np.argmax(actions_values, axis=1)\n",
    "      actions.resize((len(decision_steps), 1))\n",
    "      # Store the action that was picked, it will be put in the trajectory later\n",
    "      for agent_index, agent_id in enumerate(decision_steps.agent_id):\n",
    "        dict_last_action_from_agent[agent_id] = actions[agent_index]\n",
    "\n",
    "      # Set the actions in the environment\n",
    "      # Unity Environments expect ActionTuple instances.\n",
    "      action_tuple = ActionTuple()\n",
    "      action_tuple.add_discrete(actions)\n",
    "      env.set_actions(behavior_name, action_tuple)\n",
    "      # Perform a step in the simulation\n",
    "      env.step()\n",
    "    return buffer, np.mean(cumulative_rewards)\n",
    "\n",
    "  @staticmethod\n",
    "  def update_q_net(\n",
    "    q_net: VisualQNetwork,\n",
    "    optimizer: torch.optim,\n",
    "    buffer: Buffer,\n",
    "    action_size: int\n",
    "  ):\n",
    "    \"\"\"\n",
    "    Performs an update of the Q-Network using the provided optimizer and buffer\n",
    "    \"\"\"\n",
    "    BATCH_SIZE = 1000\n",
    "    NUM_EPOCH = 3\n",
    "    GAMMA = 0.9\n",
    "    batch_size = min(len(buffer), BATCH_SIZE)\n",
    "    random.shuffle(buffer)\n",
    "    # Split the buffer into batches\n",
    "    batches = [\n",
    "      buffer[batch_size * start : batch_size * (start + 1)]\n",
    "      for start in range(int(len(buffer) / batch_size))\n",
    "    ]\n",
    "    for _ in range(NUM_EPOCH):\n",
    "      for batch in batches:\n",
    "        # Create the Tensors that will be fed in the network\n",
    "        obs = torch.from_numpy(np.stack([ex.obs for ex in batch]))\n",
    "        reward = torch.from_numpy(\n",
    "          np.array([ex.reward for ex in batch], dtype=np.float32).reshape(-1, 1)\n",
    "        )\n",
    "        done = torch.from_numpy(\n",
    "          np.array([ex.done for ex in batch], dtype=np.float32).reshape(-1, 1)\n",
    "        )\n",
    "        action = torch.from_numpy(np.stack([ex.action for ex in batch]))\n",
    "        next_obs = torch.from_numpy(np.stack([ex.next_obs for ex in batch]))\n",
    "\n",
    "        # Use the Bellman equation to update the Q-Network\n",
    "        target = (\n",
    "          reward\n",
    "          + (1.0 - done)\n",
    "          * GAMMA\n",
    "          * torch.max(q_net(next_obs).detach(), dim=1, keepdim=True).values\n",
    "        )\n",
    "        mask = torch.zeros((len(batch), action_size))\n",
    "        mask.scatter_(1, action, 1)\n",
    "        prediction = torch.sum(q_net(obs) * mask, dim=1, keepdim=True)\n",
    "        criterion = torch.nn.MSELoss()\n",
    "        loss = criterion(prediction, target)\n",
    "\n",
    "        # Perform the backpropagation\n",
    "        optimizer.zero_grad()\n",
    "        loss.backward()\n",
    "        optimizer.step()\n"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "vcU4ZMAEWCvX"
   },
   "source": [
    "### Run Training"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "_lIHijQfbYjh"
   },
   "source": [
    "# -----------------\n",
    "# This code is used to close an env that might not have been closed before\n",
    "try:\n",
    "  env.close()\n",
    "except:\n",
    "  pass\n",
    "# -----------------\n",
    "\n",
    "from mlagents_envs.registry import default_registry\n",
    "from mlagents_envs.environment import UnityEnvironment\n",
    "import matplotlib.pyplot as plt\n",
    "%matplotlib inline\n",
    "\n",
    "# Create the GridWorld Environment from the registry\n",
    "env = default_registry[\"GridWorld\"].make()\n",
    "print(\"GridWorld environment created.\")\n",
    "\n",
    "num_actions = 5\n",
    "\n",
    "try:\n",
    "  # Create a new Q-Network.\n",
    "  qnet = VisualQNetwork((3, 64, 84), 126, num_actions)\n",
    "\n",
    "  experiences: Buffer = []\n",
    "  optim = torch.optim.Adam(qnet.parameters(), lr= 0.001)\n",
    "\n",
    "  cumulative_rewards: List[float] = []\n",
    "\n",
    "  # The number of training steps that will be performed\n",
    "  NUM_TRAINING_STEPS = int(os.getenv('QLEARNING_NUM_TRAINING_STEPS', 70))\n",
    "  # The number of experiences to collect per training step\n",
    "  NUM_NEW_EXP = int(os.getenv('QLEARNING_NUM_NEW_EXP', 1000))\n",
    "  # The maximum size of the Buffer\n",
    "  BUFFER_SIZE = int(os.getenv('QLEARNING_BUFFER_SIZE', 10000))\n",
    "\n",
    "  for n in range(NUM_TRAINING_STEPS):\n",
    "    new_exp,_ = Trainer.generate_trajectories(env, qnet, NUM_NEW_EXP, epsilon=0.1)\n",
    "    random.shuffle(experiences)\n",
    "    if len(experiences) > BUFFER_SIZE:\n",
    "      experiences = experiences[:BUFFER_SIZE]\n",
    "    experiences.extend(new_exp)\n",
    "    Trainer.update_q_net(qnet, optim, experiences, num_actions)\n",
    "    _, rewards = Trainer.generate_trajectories(env, qnet, 100, epsilon=0)\n",
    "    cumulative_rewards.append(rewards)\n",
    "    print(\"Training step \", n+1, \"\\treward \", rewards)\n",
    "except KeyboardInterrupt:\n",
    "  print(\"\\nTraining interrupted, continue to next cell to save to save the model.\")\n",
    "finally:\n",
    "  env.close()\n",
    "\n",
    "# Show the training graph\n",
    "try:\n",
    "  plt.plot(range(NUM_TRAINING_STEPS), cumulative_rewards)\n",
    "except ValueError:\n",
    "  print(\"\\nPlot failed on interrupted training.\")\n"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "source": [
    "### Export PyTorch Model to ONNX"
   ],
   "metadata": {
    "id": "pFEP9JT5KuO-"
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "The following cell provides an example of some of the extra tensors a model needs to work for ML-Agents inference with Barracuda. The GridWorldColab scene is configured to work with this ONNX file.\n",
    "Only policy models need to be exported for inference and they need the following additional tensors:\n",
    "\n",
    "*   All models need version_number\n",
    "*   All models need memory_size\n",
    "*   Models with continuous outputs need continuous_action_output_shape\n",
    "*   Models with discrete outputs need discrete_action_output_shape and an additional mask input that matches the shape of the discrete outputs\n",
    "*   The mask input must be connected to the outputs or it will be pruned on export, if mask values aren't being set they will be 1, so multiplying the discrete outputs by the mask will have no effect"
   ],
   "metadata": {
    "id": "V0V7bH4mOnUF"
   }
  },
  {
   "cell_type": "code",
   "source": [
    "class WrapperNet(torch.nn.Module):\n",
    "    def __init__(\n",
    "            self,\n",
    "            qnet: VisualQNetwork,\n",
    "            discrete_output_sizes: List[int],\n",
    "    ):\n",
    "        \"\"\"\n",
    "        Wraps the VisualQNetwork adding extra constants and dummy mask inputs\n",
    "        required by runtime inference with Barracuda.\n",
    "\n",
    "        For environment continuous actions outputs would need to add them\n",
    "        similarly to how discrete action outputs work, both in the wrapper\n",
    "        and in the ONNX output_names / dynamic_axes.\n",
    "        \"\"\"\n",
    "        super(WrapperNet, self).__init__()\n",
    "        self.qnet = qnet\n",
    "\n",
    "        # version_number\n",
    "        #   MLAgents1_0 = 2   (not covered by this example)\n",
    "        #   MLAgents2_0 = 3\n",
    "        version_number = torch.Tensor([3])\n",
    "        self.version_number = Parameter(version_number, requires_grad=False)\n",
    "\n",
    "        # memory_size\n",
    "        # TODO: document case where memory is not zero.\n",
    "        memory_size = torch.Tensor([0])\n",
    "        self.memory_size = Parameter(memory_size, requires_grad=False)\n",
    "\n",
    "        # discrete_action_output_shape\n",
    "        output_shape=torch.Tensor([discrete_output_sizes])\n",
    "        self.discrete_shape = Parameter(output_shape, requires_grad=False)\n",
    "\n",
    "\n",
    "    # if you have discrete actions ML-agents expects corresponding a mask\n",
    "    # tensor with the same shape to exist as input\n",
    "    def forward(self, visual_obs: torch.tensor, mask: torch.tensor):\n",
    "        qnet_result = self.qnet(visual_obs)\n",
    "        # Connect mask to keep it from getting pruned\n",
    "        # Mask values will be 1 if you never call SetActionMask() in\n",
    "        # WriteDiscreteActionMask()\n",
    "        qnet_result = torch.mul(qnet_result, mask)\n",
    "        action = torch.argmax(qnet_result, dim=1, keepdim=True)\n",
    "        return [action], self.discrete_shape, self.version_number, self.memory_size\n",
    "\n",
    "\n",
    "torch.onnx.export(\n",
    "    WrapperNet(qnet, [num_actions]),\n",
    "    # A tuple with an example of the input tensors\n",
    "    (torch.tensor([experiences[0].obs]), torch.ones(1, num_actions)),\n",
    "    'GridWorldColab.onnx',\n",
    "    opset_version=9,\n",
    "    # input_names must correspond to the WrapperNet forward parameters\n",
    "    # obs will be obs_0, obs_1, etc.\n",
    "    input_names=[\"obs_0\", \"action_masks\"],\n",
    "    # output_names must correspond to the return tuple of the WrapperNet\n",
    "    # forward function.\n",
    "    output_names=[\"discrete_actions\", \"discrete_action_output_shape\",\n",
    "                  \"version_number\", \"memory_size\"],\n",
    "    # All inputs and outputs should have their 0th dimension be designated\n",
    "    # as 'batch'\n",
    "    dynamic_axes={'obs_0': {0: 'batch'},\n",
    "                  'action_masks': {0: 'batch'},\n",
    "                  'discrete_actions': {0: 'batch'},\n",
    "                  'discrete_action_output_shape': {0: 'batch'}\n",
    "                 }\n",
    "    )"
   ],
   "metadata": {
    "id": "T0zzg8fWpTYO"
   },
   "execution_count": null,
   "outputs": []
  }
 ]
}