examples/quickstart.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "ULdrhOaVbsdO"
      },
      "source": [
        "# Acme: Quickstart\n",
        "## Guide to installing Acme and training your first D4PG agent.\n",
        "# <a href=\"https://colab.research.google.com/github/deepmind/acme/blob/master/examples/quickstart.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "ogw2P040-F5P"
      },
      "source": [
        "## Select your environment library\n",
        "Note: `dm_control` requires a valid Mujoco license."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "cellView": "form",
        "colab": {},
        "colab_type": "code",
        "id": "xflSXJPS8Qpm"
      },
      "outputs": [],
      "source": [
        "environment_library = 'gym'  # @param ['dm_control', 'gym']"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "U0UfMBOL_T2l"
      },
      "source": [
        "## Add your Mujoco license here\n",
        "Note: only required for `dm_control`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "QbCY7t4--eOd"
      },
      "outputs": [],
      "source": [
        "mjkey = \"\"\"\n",
        "\"\"\".strip()\n",
        "\n",
        "if not mjkey and environment_library == 'dm_control':\n",
        "  raise ValueError(\n",
        "      'A Mujoco license is required for `dm_control`, if you do not have on '\n",
        "      'consider selecting `gym` from the dropdown menu in the cell above.')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "xaJxoatMhJ71"
      },
      "source": [
        "## Installation"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "ovuCuHCC78Zu"
      },
      "source": [
        "### Install Acme"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "cellView": "both",
        "colab": {},
        "colab_type": "code",
        "id": "KH3O0zcXUeun"
      },
      "outputs": [],
      "source": [
        "!pip install dm-acme\n",
        "!pip install dm-acme[reverb]\n",
        "!pip install dm-acme[tf]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "VEEj3Qw60y73"
      },
      "source": [
        "### Install the environment library\n",
        "\n",
        "Without a valid license you won't be able to use the `dm_control` environments but can still follow this colab using the `gym` environments.\n",
        "\n",
        "If you have a personal Mujoco license (_not_ an institutional one), you may \n",
        "need to follow the instructions at https://research.google.com/colaboratory/local-runtimes.html to run a Jupyter kernel on your local machine.\n",
        "This will allow you to install `dm_control` by following instructions in\n",
        "https://github.com/deepmind/dm_control and using a personal MuJoCo license.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "cellView": "both",
        "colab": {},
        "colab_type": "code",
        "id": "IbZxYDxzoz5R"
      },
      "outputs": [],
      "source": [
        "#@test {\"skip\": true}\n",
        "if environment_library == 'dm_control':\n",
        "  mujoco_dir = \"$HOME/.mujoco\"\n",
        "\n",
        "  # Install OpenGL dependencies\n",
        "  !apt-get update && apt-get install -y --no-install-recommends \\\n",
        "    libgl1-mesa-glx libosmesa6 libglew2.0\n",
        "\n",
        "  # Get MuJoCo binaries\n",
        "  !wget -q https://www.roboti.us/download/mujoco200_linux.zip -O mujoco.zip\n",
        "  !unzip -o -q mujoco.zip -d \"$mujoco_dir\"\n",
        "\n",
        "  # Copy over MuJoCo license\n",
        "  !echo \"$mjkey\" > \"$mujoco_dir/mjkey.txt\"\n",
        "\n",
        "  # Install dm_control\n",
        "  !pip install dm_control\n",
        "\n",
        "  # Configure dm_control to use the OSMesa rendering backend\n",
        "  %env MUJOCO_GL=osmesa\n",
        "\n",
        "  # Check that the installation succeeded\n",
        "  try:\n",
        "    from dm_control import suite\n",
        "    env = suite.load('cartpole', 'swingup')\n",
        "    pixels = env.physics.render()\n",
        "  except Exception as e:\n",
        "    raise e from RuntimeError(\n",
        "        'Something went wrong during installation. Check the shell output above '\n",
        "        'for more information. If you do not have a valid Mujoco license, '\n",
        "        'consider selecting `gym` in the dropdown menu at the top of this Colab.')\n",
        "  else:\n",
        "    del suite, env, pixels\n",
        "\n",
        "elif environment_library == 'gym':\n",
        "  !pip install gym"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "Cl8eyWblSs-z"
      },
      "source": [
        "### Install visualization packages"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "aSM7KHDFSsQS"
      },
      "outputs": [],
      "source": [
        "!sudo apt-get install -y xvfb ffmpeg\n",
        "!pip install imageio\n",
        "!pip install PILLOW\n",
        "!pip install pyvirtualdisplay"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "c-H2d6UZi7Sf"
      },
      "source": [
        "## Import Modules"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "cellView": "both",
        "colab": {},
        "colab_type": "code",
        "id": "HJ74Id-8MERq"
      },
      "outputs": [],
      "source": [
        "import IPython\n",
        "\n",
        "from acme import environment_loop\n",
        "from acme import specs\n",
        "from acme import wrappers\n",
        "from acme.agents.tf import d4pg\n",
        "from acme.tf import networks\n",
        "from acme.tf import utils as tf2_utils\n",
        "from acme.utils import loggers\n",
        "import numpy as np\n",
        "import sonnet as snt\n",
        "\n",
        "# Import the selected environment lib\n",
        "if environment_library == 'dm_control':\n",
        "  from dm_control import suite\n",
        "elif environment_library == 'gym':\n",
        "  import gym\n",
        "\n",
        "# Imports required for visualization\n",
        "import pyvirtualdisplay\n",
        "import imageio\n",
        "import base64\n",
        "\n",
        "# Set up a virtual display for rendering.\n",
        "display = pyvirtualdisplay.Display(visible=0, size=(1400, 900)).start()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "I6KuVGSk4uc9"
      },
      "source": [
        "## Load an environment\n",
        "\n",
        "We can now load an environment. In what follows we'll create an environment and grab the environment's specifications."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "cellView": "both",
        "colab": {},
        "colab_type": "code",
        "id": "4PVlHtGF5yzt"
      },
      "outputs": [],
      "source": [
        "if environment_library == 'dm_control':\n",
        "  environment = suite.load('cartpole', 'balance')\n",
        "  \n",
        "elif environment_library == 'gym':\n",
        "  environment = gym.make('MountainCarContinuous-v0')\n",
        "  environment = wrappers.GymWrapper(environment)  # To dm_env interface.\n",
        "\n",
        "else:\n",
        "  raise ValueError(\n",
        "      \"Unknown environment library: {};\".format(environment_name) +\n",
        "      \"choose among ['dm_control', 'gym'].\")\n",
        "\n",
        "# Make sure the environment outputs single-precision floats.\n",
        "environment = wrappers.SinglePrecisionWrapper(environment)\n",
        "\n",
        "# Grab the spec of the environment.\n",
        "environment_spec = specs.make_environment_spec(environment)\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "BukOfOsmtSQn"
      },
      "source": [
        " ## Create a D4PG agent"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "3Jcjk1w6oHVX"
      },
      "outputs": [],
      "source": [
        "#@title Build agent networks\n",
        "\n",
        "# Get total number of action dimensions from action spec.\n",
        "num_dimensions = np.prod(environment_spec.actions.shape, dtype=int)\n",
        "\n",
        "# Create the shared observation network; here simply a state-less operation.\n",
        "observation_network = tf2_utils.batch_concat\n",
        "\n",
        "# Create the deterministic policy network.\n",
        "policy_network = snt.Sequential([\n",
        "    networks.LayerNormMLP((256, 256, 256), activate_final=True),\n",
        "    networks.NearZeroInitializedLinear(num_dimensions),\n",
        "    networks.TanhToSpec(environment_spec.actions),\n",
        "])\n",
        "\n",
        "# Create the distributional critic network.\n",
        "critic_network = snt.Sequential([\n",
        "    # The multiplexer concatenates the observations/actions.\n",
        "    networks.CriticMultiplexer(),\n",
        "    networks.LayerNormMLP((512, 512, 256), activate_final=True),\n",
        "    networks.DiscreteValuedHead(vmin=-150., vmax=150., num_atoms=51),\n",
        "])\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "9CD2sNK-oA9S"
      },
      "outputs": [],
      "source": [
        "# Create a logger for the agent and environment loop.\n",
        "agent_logger = loggers.TerminalLogger(label='agent', time_delta=10.)\n",
        "env_loop_logger = loggers.TerminalLogger(label='env_loop', time_delta=10.)\n",
        "\n",
        "# Create the D4PG agent.\n",
        "agent = d4pg.D4PG(\n",
        "    environment_spec=environment_spec,\n",
        "    policy_network=policy_network,\n",
        "    critic_network=critic_network,\n",
        "    observation_network=observation_network,\n",
        "    sigma=1.0,\n",
        "    logger=agent_logger,\n",
        "    checkpoint=False\n",
        ")\n",
        "\n",
        "# Create an loop connecting this agent to the environment created above.\n",
        "env_loop = environment_loop.EnvironmentLoop(\n",
        "    environment, agent, logger=env_loop_logger)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "oKeGQxzitXYC"
      },
      "source": [
        "## Run a training loop"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "VWZd5N-Qoz82"
      },
      "outputs": [],
      "source": [
        "# Run a `num_episodes` training episodes.\n",
        "# Rerun this cell until the agent has learned the given task.\n",
        "env_loop.run(num_episodes=100)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "Do57Ql4ZsWDu"
      },
      "source": [
        "## Visualize an evaluation loop\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "BJXkfg6LSZ-h"
      },
      "source": [
        "### Helper functions for rendering and vizualization"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "OIJRbtAlxQVu"
      },
      "outputs": [],
      "source": [
        "# Create a simple helper function to render a frame from the current state of\n",
        "# the environment.\n",
        "if environment_library == 'dm_control':\n",
        "  def render(env):\n",
        "    return env.physics.render(camera_id=0)\n",
        "elif environment_library == 'gym':\n",
        "  def render(env):\n",
        "    return env.environment.render(mode='rgb_array')\n",
        "else:\n",
        "  raise ValueError(\n",
        "      \"Unknown environment library: {};\".format(environment_name) +\n",
        "      \"choose among ['dm_control', 'gym'].\")\n",
        "\n",
        "def display_video(frames, filename='temp.mp4'):\n",
        "  \"\"\"Save and display video.\"\"\"\n",
        "\n",
        "  # Write video\n",
        "  with imageio.get_writer(filename, fps=60) as video:\n",
        "    for frame in frames:\n",
        "      video.append_data(frame)\n",
        "\n",
        "  # Read video and display the video\n",
        "  video = open(filename, 'rb').read()\n",
        "  b64_video = base64.b64encode(video)\n",
        "  video_tag = ('<video  width=\"320\" height=\"240\" controls alt=\"test\" '\n",
        "               'src=\"data:video/mp4;base64,{0}\">').format(b64_video.decode())\n",
        "\n",
        "  return IPython.display.HTML(video_tag)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "Rh8x46UKSf2c"
      },
      "source": [
        "### Run and visualize the agent in the environment for an episode"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "2__mFiraWND1"
      },
      "outputs": [],
      "source": [
        "timestep = environment.reset()\n",
        "frames = [render(environment)]\n",
        "\n",
        "while not timestep.last():\n",
        "  # Simple environment loop.\n",
        "  action = agent.select_action(timestep.observation)\n",
        "  timestep = environment.step(action)\n",
        "\n",
        "  # Render the scene and add it to the frame stack.\n",
        "  frames.append(render(environment))\n",
        "\n",
        "# Save and display a video of the behaviour.\n",
        "display_video(np.array(frames))"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "collapsed_sections": [],
      "name": "Acme: Quickstart",
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}