{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# ML-Agents PettingZoo Wrapper"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "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",
    "  pro_bar.update(progress(5, 100))\n",
    "  !apt-get install daemon >/dev/null 2>&1\n",
    "  pro_bar.update(progress(10, 100))\n",
    "  !apt-get install wget >/dev/null 2>&1\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 >/dev/null 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 >/dev/null 2>&1\n",
    "  pro_bar.update(progress(40, 100))\n",
    "  !dpkg -i libxfont1_1.5.1-1ubuntu0.16.04.4_amd64.deb >/dev/null 2>&1\n",
    "  pro_bar.update(progress(50, 100))\n",
    "  !dpkg -i xvfb.deb >/dev/null 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))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Installing ml-agents"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "try:\n",
    "  import mlagents\n",
    "  print(\"ml-agents already installed\")\n",
    "except ImportError:\n",
    "  !git clone -b main --single-branch https://github.com/Unity-Technologies/ml-agents.git\n",
    "  !python -m pip install -q ./ml-agents/ml-agents-envs\n",
    "  !python -m pip install -q ./ml-agents/ml-agents\n",
    "  print(\"Installed ml-agents\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Run the Environment"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true,
    "tags": []
   },
   "source": [
    "List of available environments:\n",
    "* Basic\n",
    "* ThreeDBall\n",
    "* ThreeDBallHard\n",
    "* GridWorld\n",
    "* Hallway\n",
    "* VisualHallway\n",
    "* CrawlerDynamicTarget\n",
    "* CrawlerStaticTarget\n",
    "* Bouncer\n",
    "* SoccerTwos\n",
    "* PushBlock\n",
    "* VisualPushBlock\n",
    "* WallJump\n",
    "* Tennis\n",
    "* Reacher\n",
    "* Pyramids\n",
    "* VisualPyramids\n",
    "* Walker\n",
    "* FoodCollector\n",
    "* VisualFoodCollector\n",
    "* StrikersVsGoalie\n",
    "* WormStaticTarget\n",
    "* WormDynamicTarget"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Start Environment with PettingZoo Wrapper"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "YSf-WhxbqtLw"
   },
   "outputs": [],
   "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",
    "import numpy as np\n",
    "from mlagents_envs.envs import StrikersVsGoalie # import unity environment\n",
    "env = StrikersVsGoalie.env()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Stepping the environment\n",
    "\n",
    "Example of interacting with the environment in basic RL loop. It follows the same interface as described in [PettingZoo API page](https://www.pettingzoo.ml/api)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "dhtl0mpeqxYi"
   },
   "outputs": [],
   "source": [
    "num_cycles = 10\n",
    "\n",
    "env.reset()\n",
    "for agent in env.agent_iter(env.num_agents * num_cycles):\n",
    "    prev_observe, reward, done, info = env.last()\n",
    "    if isinstance(prev_observe, dict) and 'action_mask' in prev_observe:\n",
    "        action_mask = prev_observe['action_mask']\n",
    "    if done:\n",
    "        action = None\n",
    "    else:\n",
    "        action = env.action_spaces[agent].sample() # randomly choose an action for example\n",
    "    env.step(action)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Additional Environment API\n",
    "\n",
    "All the API described in the `Additional Environment API` section in the [PettingZoo API page](https://www.pettingzoo.ml/api) are all supported. A few examples are shown below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# `agents`: a list of the names of all current agents\n",
    "print(\"Agent names:\", env.agents)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# `agent_selection`: the currently agent that an action can be taken for.\n",
    "print(\"Current agent:\", env.agent_selection)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# `observation_spaces`: a dict of the observation spaces of every agent, keyed by name.\n",
    "print(\"Observation space of current agent:\", env.observation_spaces[env.agent_selection])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# `action_spaces`: a dict of the observation spaces of every agent, keyed by name.\n",
    "print(\"Action space of current agent:\", env.action_spaces[env.agent_selection])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Close the Environment to free the port it is using"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "a7KatdThq7OV"
   },
   "outputs": [],
   "source": [
    "env.close()"
   ]
  }
 ],
 "metadata": {
  "colab": {
   "collapsed_sections": [],
   "name": "Colab-UnityEnvironment-1-Run.ipynb",
   "private_outputs": true,
   "provenance": [],
   "toc_visible": true
  },
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}