dqn_zoo_plots.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "WpaLpc1tNBI2"
      },
      "source": [
        "# DQN Zoo plots"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "oCbIAwKAefTr"
      },
      "source": [
        "Copyright 2021 DeepMind Technologies Limited\n",
        "\n",
        "Licensed under the Apache License, Version 2.0 (the \"License\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "fZfgzl9YeEs0"
      },
      "outputs": [],
      "source": [
        "# @title Licence {form-width: '30%'}\n",
        "\n",
        "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "#     https://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "NasstDUaSu8V"
      },
      "outputs": [],
      "source": [
        "import os.path\n",
        "\n",
        "import matplotlib.pyplot as plt\n",
        "from matplotlib import ticker\n",
        "import numpy as np\n",
        "import pandas as pd\n",
        "import seaborn as sns"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "xJaoxmbLMmjp"
      },
      "outputs": [],
      "source": [
        "# @title Defaults {run:\"auto\", form-width: '30%'}\n",
        "\n",
        "backend_format = \"retina\"  # @param [\"retina\", \"\"]\n",
        "%config InlineBackend.figure_format = backend_format\n",
        "\n",
        "sns.set_context(\"notebook\")\n",
        "sns.set_style(\"ticks\")\n",
        "\n",
        "plt.rcParams['ytick.right'] = True\n",
        "plt.rcParams['axes.autolimit_mode'] = 'round_numbers'\n",
        "plt.rcParams['axes.xmargin'] = 0\n",
        "plt.rcParams['axes.ymargin'] = 0\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "RGzt_5KNLb1l"
      },
      "outputs": [],
      "source": [
        "!wget https://github.com/deepmind/dqn_zoo/raw/master/results.tar.gz"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "78lHkC4OMOTv"
      },
      "outputs": [],
      "source": [
        "!tar -xvf results.tar.gz"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "GzPJrxwbNJys"
      },
      "outputs": [],
      "source": [
        "experiments = [\n",
        "    dict(\n",
        "        agent_id='dqn',\n",
        "        agent_name='DQN',\n",
        "        color='grey'),\n",
        "    dict(\n",
        "        agent_id='double_q',\n",
        "        agent_name='Double DQN',\n",
        "        color='purple'),\n",
        "    dict(\n",
        "        agent_id='prioritized',\n",
        "        agent_name='Prioritized DQN',\n",
        "        color='blue'),\n",
        "    dict(\n",
        "        agent_id='c51',\n",
        "        agent_name='C51',\n",
        "        color='orange'),\n",
        "    dict(\n",
        "        agent_id='qrdqn',\n",
        "        agent_name='QR-DQN',\n",
        "        color='red'),\n",
        "    dict(\n",
        "        agent_id='rainbow',\n",
        "        agent_name='Rainbow',\n",
        "        color='black'),\n",
        "    dict(\n",
        "        agent_id='iqn',\n",
        "        agent_name='IQN',\n",
        "        color='green'),\n",
        "]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "hWF8jhy4SWBj"
      },
      "outputs": [],
      "source": [
        "# @title Game names {'form-width': '30%'}\n",
        "\n",
        "GAME_NAMES = [\n",
        "    ('alien', 'Alien'),\n",
        "    ('amidar', 'Amidar'),\n",
        "    ('assault', 'Assault'),\n",
        "    ('asterix', 'Asterix'),\n",
        "    ('asteroids', 'Asteroids'),\n",
        "    ('atlantis', 'Atlantis'),\n",
        "    ('bank_heist', 'Bank Heist'),\n",
        "    ('battle_zone', 'Battlezone'),\n",
        "    ('beam_rider', 'Beam Rider'),\n",
        "    ('berzerk', 'Berzerk'),\n",
        "    ('bowling', 'Bowling'),\n",
        "    ('boxing', 'Boxing'),\n",
        "    ('breakout', 'Breakout'),\n",
        "    ('centipede', 'Centipede'),\n",
        "    ('chopper_command', 'Chopper Command'),\n",
        "    ('crazy_climber', 'Crazy Climber'),\n",
        "    ('defender', 'Defender'),\n",
        "    ('demon_attack', 'Demon Attack'),\n",
        "    ('double_dunk', 'Double Dunk'),\n",
        "    ('enduro', 'Enduro'),\n",
        "    ('fishing_derby', 'Fishing Derby'),\n",
        "    ('freeway', 'Freeway'),\n",
        "    ('frostbite', 'Frostbite'),\n",
        "    ('gopher', 'Gopher'),\n",
        "    ('gravitar', 'Gravitar'),\n",
        "    ('hero', 'H.E.R.O.'),\n",
        "    ('ice_hockey', 'Ice Hockey'),\n",
        "    ('jamesbond', 'James Bond 007'),\n",
        "    ('kangaroo', 'Kangaroo'),\n",
        "    ('krull', 'Krull'),\n",
        "    ('kung_fu_master', 'Kung-Fu Master'),\n",
        "    ('montezuma_revenge', 'Montezuma’s Revenge'),\n",
        "    ('ms_pacman', 'Ms. Pac-Man'),\n",
        "    ('name_this_game', 'Name This Game'),\n",
        "    ('phoenix', 'Phoenix'),\n",
        "    ('pitfall', 'Pitfall!'),\n",
        "    ('pong', 'Pong'),\n",
        "    ('private_eye', 'Private Eye'),\n",
        "    ('qbert', 'Q*bert'),\n",
        "    ('riverraid', 'River Raid'),\n",
        "    ('road_runner', 'Road Runner'),\n",
        "    ('robotank', 'Robotank'),\n",
        "    ('seaquest', 'Seaquest'),\n",
        "    ('skiing', 'Skiing'),\n",
        "    ('solaris', 'Solaris'),\n",
        "    ('space_invaders', 'Space Invaders'),\n",
        "    ('star_gunner', 'Stargunner'),\n",
        "    ('surround', 'Surround'),\n",
        "    ('tennis', 'Tennis'),\n",
        "    ('time_pilot', 'Time Pilot'),\n",
        "    ('tutankham', 'Tutankham'),\n",
        "    ('up_n_down', 'Up’n Down'),\n",
        "    ('venture', 'Venture'),\n",
        "    ('video_pinball', 'Video Pinball'),\n",
        "    ('wizard_of_wor', 'Wizard of Wor'),\n",
        "    ('yars_revenge', 'Yars’ Revenge'),\n",
        "    ('zaxxon', 'Zaxxon'),\n",
        "]\n",
        "GAME_NAME_MAP = dict(GAME_NAMES)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "3UGU6KYsSa2D"
      },
      "outputs": [],
      "source": [
        "def load_experiment_data_from_results_csv_dir(experiment_details, results_dir):\n",
        "  df_exps = []\n",
        "  for ed in experiment_details:\n",
        "    csv_file = os.path.join(results_dir, ed['agent_id'] + '.csv')\n",
        "    with open(csv_file, 'r') as f:\n",
        "      df = pd.read_csv(f, index_col=0)\n",
        "    df = df.assign(agent_id=ed['agent_id'], agent_name=ed['agent_name'])\n",
        "    df_exps.append(df)\n",
        "\n",
        "  df_exp = pd.concat(df_exps, sort=True).reset_index(drop=True)\n",
        "  return df_exp\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "iv-0k6DLUpoC"
      },
      "outputs": [],
      "source": [
        "df_exp_raw = load_experiment_data_from_results_csv_dir(experiments, '.')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "S2dVyg6NTBvf"
      },
      "outputs": [],
      "source": [
        "def moving_average(values, window_size):\n",
        "  # numpy.convolve uses zero for initial missing values, so is not suitable.\n",
        "  numerator = np.nancumsum(values)\n",
        "  # The sum of the last window_size values.\n",
        "  numerator[window_size:] = numerator[window_size:] - numerator[:-window_size]\n",
        "  denominator = np.ones(len(values)) * window_size\n",
        "  denominator[:window_size] = np.arange(1, window_size + 1)\n",
        "  smoothed = numerator / denominator\n",
        "  assert values.shape == smoothed.shape\n",
        "  return smoothed\n",
        "\n",
        "\n",
        "def smooth(df, smoothing_window, index_columns, columns):\n",
        "  dfg = df.groupby(index_columns)\n",
        "  for col in columns:\n",
        "    df[col] = dfg[col].transform(\n",
        "        lambda s: moving_average(s.values, smoothing_window))\n",
        "  return df\n",
        "\n",
        "\n",
        "def environment_pretty(row):\n",
        "  return GAME_NAME_MAP[row['environment_name']]\n",
        "\n",
        "\n",
        "def add_columns(df):\n",
        "  df['environment_pretty'] = df.apply(environment_pretty, axis=1)\n",
        "  df['frame_millions'] = df['frame'] // int(1e6)\n",
        "  return df\n",
        "\n",
        "\n",
        "def smooth_dataframe(df):\n",
        "  return smooth(\n",
        "      df,\n",
        "      smoothing_window=10,\n",
        "      index_columns=['agent_id', 'environment_name', 'seed'],\n",
        "      columns=[\n",
        "          'eval_frame_rate',\n",
        "          'train_frame_rate',\n",
        "          'eval_episode_return',\n",
        "          'train_episode_return',\n",
        "          'normalized_return',\n",
        "      ])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "HL3bj_BpTLrw"
      },
      "outputs": [],
      "source": [
        "df_exp = df_exp_raw.pipe(add_columns).pipe(smooth_dataframe)\n",
        "df = df_exp.sort_values(by=['agent_id', 'environment_name', 'seed', 'frame'])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "4IlnXjDaTz14"
      },
      "outputs": [],
      "source": [
        "def make_agent_hue_kws(experiments):\n",
        "  pairs = [(exp['agent_name'], exp['color']) for exp in experiments]\n",
        "  agent_names, colors = zip(*pairs)\n",
        "  hue_kws = dict(color=colors)\n",
        "  return list(agent_names), hue_kws\n",
        "\n",
        "\n",
        "def plot_individual(df, agent_names, hue_kws):\n",
        "  g = sns.FacetGrid(\n",
        "      df.query('agent_name == %s' % agent_names),\n",
        "      row=None,\n",
        "      col='environment_pretty',\n",
        "      hue='agent_name',\n",
        "      height=2.5,\n",
        "      aspect=1.5,\n",
        "      col_wrap=3,\n",
        "      hue_order=agent_names,\n",
        "      sharey=False,\n",
        "      hue_kws=hue_kws,\n",
        "  )\n",
        "\n",
        "  g.map(\n",
        "      sns.lineplot,\n",
        "      'frame_millions',\n",
        "      'eval_episode_return',\n",
        "      estimator=np.median,\n",
        "      ci=100,\n",
        "      alpha=0.5,\n",
        "      linewidth=3,\n",
        "  )\n",
        "\n",
        "  g.set_titles('{col_name}')\n",
        "  g.despine(left=False, top=True, right=False, bottom=False)\n",
        "  g.set_axis_labels('Frame (millions)', '')\n",
        "\n",
        "  for ax in g.axes:\n",
        "    ax.yaxis.set_major_locator(ticker.MaxNLocator(nbins=6))\n",
        "    ax.xaxis.set_major_locator(ticker.MaxNLocator(nbins=4))\n",
        "\n",
        "  g.fig.set_facecolor('white')\n",
        "  g.fig.tight_layout()\n",
        "\n",
        "  # Create legend from the final axes.\n",
        "  legend_x_margin = 0.03\n",
        "  legend_y_offset_inches = -0.4\n",
        "  legend_y_offset = legend_y_offset_inches / g.fig.get_figheight()\n",
        "  g.axes[-1].legend(\n",
        "      bbox_to_anchor=(\n",
        "          legend_x_margin,\n",
        "          legend_y_offset,\n",
        "          1 - 2 * legend_x_margin,\n",
        "          0,\n",
        "      ),\n",
        "      bbox_transform=g.fig.transFigure,\n",
        "      mode='expand',\n",
        "      ncol=10,\n",
        "      borderaxespad=0,\n",
        "      loc='lower center',\n",
        "      frameon=False,\n",
        "  )\n",
        "  return g\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "TaDezDLzUBqK"
      },
      "outputs": [],
      "source": [
        "# This can take a very long time to run.\n",
        "plot_individual(df, *make_agent_hue_kws(experiments))\n",
        "pass"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "eKhLQ95CTa9g"
      },
      "outputs": [],
      "source": [
        "def plot_summary(df, agent_names, hue_kws):\n",
        "  df = df.query('agent_name == %s' % agent_names)\n",
        "  df = df.groupby(['agent_name', 'frame', 'seed']).median().reset_index()\n",
        "  df['normalized_return_percent'] = df['normalized_return'] * 100\n",
        "\n",
        "  fig, ax = plt.subplots(figsize=(9, 6))\n",
        "  sns.lineplot(\n",
        "      x=df['frame_millions'],\n",
        "      y=df['normalized_return_percent'],\n",
        "      hue=df['agent_name'],\n",
        "      estimator=np.median,\n",
        "      ci=100,\n",
        "      alpha=0.5,\n",
        "      linewidth=3,\n",
        "      hue_order=agent_names,\n",
        "      palette=hue_kws['color'],\n",
        "      legend='brief',\n",
        "      ax=ax,\n",
        "  )\n",
        "  ax.set_title('')\n",
        "  ax.spines['right'].set_visible(True)\n",
        "  ax.spines['top'].set_visible(False)\n",
        "  ax.spines['left'].set_visible(True)\n",
        "  ax.spines['bottom'].set_visible(True)\n",
        "  ax.set_xlabel('Frame (millions)')\n",
        "  ax.set_ylabel('Median human-normalized score')\n",
        "  ax.set(ylim=(0, 250))\n",
        "\n",
        "  ax.yaxis.set_major_locator(ticker.MaxNLocator(nbins=3))\n",
        "  ax.xaxis.set_major_locator(ticker.MaxNLocator(nbins=4))\n",
        "  ax.yaxis.set_major_formatter(ticker.PercentFormatter())\n",
        "\n",
        "  legend = ax.legend(loc='upper left', title='', frameon=False)\n",
        "  plt.setp(legend.get_lines(), linewidth=3, alpha=0.5)\n",
        "\n",
        "  fig.set_facecolor('white')\n",
        "  fig.tight_layout()\n",
        "  return fig"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "0F0WBQxuTrdY"
      },
      "outputs": [],
      "source": [
        "agent_id_to_experiment = {e['agent_id']: e for e in experiments}\n",
        "agent_id_ordered_by_performance = [\n",
        "    'rainbow',\n",
        "    'iqn',\n",
        "    'qrdqn',\n",
        "    'c51',\n",
        "    'prioritized',\n",
        "    'double_q',\n",
        "    'dqn',\n",
        "]\n",
        "experiments_ordered_by_performance = [\n",
        "    agent_id_to_experiment[agent_id]\n",
        "    for agent_id in agent_id_ordered_by_performance\n",
        "]\n",
        "plot_summary(df, *make_agent_hue_kws(experiments_ordered_by_performance))\n",
        "pass"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "collapsed_sections": [],
      "name": "dqn_zoo_plots.ipynb",
      "private_outputs": true,
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}