{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pickle\n",
    "import os\n",
    "\n",
    "from scipy.signal import medfilt\n",
    "from matplotlib import pyplot as plt\n",
    "import matplotlib as mpl\n",
    "import numpy as np\n",
    "import itertools \n",
    "import pandas as pd"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Learning Curve Plots"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_learning_average(env_name, ax, experiment, window_size, var = True, OnePlayer = False):\n",
    "    if OnePlayer:\n",
    "        base_dir = os.getcwd() + '/models_OnePlayer/' + env_name\n",
    "    else:\n",
    "        base_dir = os.getcwd() + '/models/' + env_name\n",
    "    results = {}\n",
    "        \n",
    "    for exp in experiment:\n",
    "        results[exp] = {}\n",
    "        for seed in range(5):\n",
    "            try:\n",
    "                if(OnePlayer):\n",
    "                    with open(base_dir + '/' + exp +'/nr_mdp_0_1/' + str(seed) + '/results', 'rb') as f:\n",
    "                        results[exp][seed] = pickle.load(f)\n",
    "                else:\n",
    "                    with open(base_dir + '/' + exp +'/nr_mdp_0.1_1/' + str(seed) + '/results', 'rb') as f:\n",
    "                        results[exp][seed] = pickle.load(f)                \n",
    "            except IOError:\n",
    "                pass\n",
    "            \n",
    "    avg, std = plot_learning_curves(results, ax, env_name, window_size, var, OnePlayer, best = False)\n",
    "    return avg, std\n",
    "\n",
    "def plot_learning_best(env_name, seed, ax, experiment, window_size, var = True, OnePlayer = False):\n",
    "    if OnePlayer:\n",
    "        base_dir = os.getcwd() + '/models_OnePlayer/' + env_name\n",
    "    else:\n",
    "        base_dir = os.getcwd() + '/models/' + env_name\n",
    "    results = {}\n",
    "        \n",
    "    for i, exp in enumerate(experiment):\n",
    "        results[exp] = {}\n",
    "        try:\n",
    "            if(OnePlayer):\n",
    "                with open(base_dir + '/' + exp +'/nr_mdp_0_1/' + str(seed[i]) + '/results', 'rb') as f:\n",
    "                    results[exp] = pickle.load(f)\n",
    "            else:\n",
    "                with open(base_dir + '/' + exp +'/nr_mdp_0.1_1/' + str(seed[i]) + '/results', 'rb') as f:\n",
    "                    results[exp] = pickle.load(f)                \n",
    "        except IOError:\n",
    "            pass\n",
    "            \n",
    "    avg, std = plot_learning_curves(results, ax, env_name, window_size, var, OnePlayer, best = True)\n",
    "    return avg, std\n",
    "   \n",
    "def plot_learning_curves(results, ax, env_name, window_size, var, OnePlayer, best):    \n",
    "    colors = ['#396ab1', '#3e9651', '#cc2529', '#396ab1', '#da7c30', '#94823d', '#535154', '#006400', '#00FF00', '#800000', '#F08080', '#FFFF00', '#000000', '#C0C0C0']\n",
    "\n",
    "    idx = 0\n",
    "    final_avg = []\n",
    "    final_std = []\n",
    "    for exp in results:\n",
    "        if(best):\n",
    "            merged = list(itertools.chain(*(results[exp]['eval_rewards'])))\n",
    "        else:\n",
    "            reward = []    \n",
    "            for seed in results[exp]:\n",
    "                reward.append(list(itertools.chain(*(results[exp][seed]['eval_rewards']))))\n",
    "            merged = list(itertools.chain.from_iterable(reward))\n",
    "        out = np.array(merged) \n",
    "        out = out.reshape(-1,2)\n",
    "        df = pd.DataFrame({'Column1':out[:,0],'Column2':out[:,1]})\n",
    "        length = int(df.Column1.values[-1]/window_size)\n",
    "        x = np.zeros(length)\n",
    "        y_avg = np.zeros(length)\n",
    "        y_std = np.zeros(length)\n",
    "        for i in range(1, length):\n",
    "            data = (df[(df.Column1 <= window_size * i) & (df.Column1 > window_size * (i-1))]).Column2.values\n",
    "            x[i] = window_size * i\n",
    "            y_avg[i] = np.mean(data)\n",
    "            y_std[i] = np.std(data)\n",
    "        avg = y_avg[-10:]\n",
    "        final_avg.append(np.mean(avg))\n",
    "        final_std.append(np.std(avg))\n",
    "        if 'RMSprop' in exp:\n",
    "            l = 'GAD(RMSprop)'\n",
    "        elif 'ExtraAdam' in exp:\n",
    "            l = 'Extra-Adam'\n",
    "        else:\n",
    "            l = 'MixedNE-LD(RMSProp)'       \n",
    "        ax.plot(x, y_avg, color = colors[idx], label = l)\n",
    "        if (var):\n",
    "            ax.fill_between(x, (y_avg)-(y_std), (y_avg)+(y_std), facecolor=colors[idx], alpha=0.4, interpolate=True)\n",
    "        idx += 1  \n",
    "    \n",
    "    ax.set_title(env_name)    \n",
    "    ax.set_xlabel('Timesteps')\n",
    "    ax.set_ylabel('Reward')\n",
    "    ax.ticklabel_format(style='scientific', axis='x', scilimits=(0,0))    \n",
    "    return final_avg, final_std"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_learning_all(experiment, title, comparison, OnePlayer, best, seed=None):\n",
    "    cols = 2\n",
    "    rows = 4\n",
    "    envs = ['Walker2d-v2', 'HalfCheetah-v2', 'Hopper-v2', 'Ant-v2', 'Swimmer-v2', 'Reacher-v2', 'Humanoid-v2', 'InvertedPendulum-v2']\n",
    "    fill = [True, False]\n",
    "    for f in fill:\n",
    "        fig, axs = plt.subplots(rows, cols, figsize=(20*cols/3, 9*rows/2))\n",
    "        for i, env in enumerate(envs):\n",
    "            ax = axs.flatten()[i]\n",
    "            if(comparison):\n",
    "                exp = experiment[i]\n",
    "            else:\n",
    "                exp = experiment\n",
    "            if(best):\n",
    "                plot_learning_best(env, seed[i], ax, exp, window_size = 3000, var = f, OnePlayer = OnePlayer)\n",
    "            else:\n",
    "                plot_learning_average(env, ax, exp, window_size = 3000, var = f, OnePlayer = OnePlayer)\n",
    "            if(comparison):\n",
    "                ax.legend()\n",
    "            else:\n",
    "                if i == 0:\n",
    "                    ax.legend()\n",
    "        plt.tight_layout()\n",
    "        \n",
    "        if(OnePlayer):\n",
    "            folder = 'OnePlayer'\n",
    "        else:\n",
    "            folder = 'TwoPlayer'\n",
    "            \n",
    "        if(f):\n",
    "            if(best):\n",
    "                plt.savefig('plots/'+ folder + '/Learning_curve/' + title +'_best_fill.pdf')\n",
    "            else:\n",
    "                plt.savefig('plots/'+ folder + '/Learning_curve/' + title +'_average_fill.pdf')                    \n",
    "        else:\n",
    "            if(best):\n",
    "                plt.savefig('plots/'+ folder + '/Learning_curve/' + title +'_best.pdf')\n",
    "            else:\n",
    "                plt.savefig('plots/'+ folder + '/Learning_curve/' + title +'_average.pdf')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 645372 characters]\n",
      "text/plain": [
       "<Figure size 960x1296 with 8 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 484996 characters]\n",
      "text/plain": [
       "<Figure size 960x1296 with 8 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "## Comparison Two Player Average\n",
    "Walker_exp = ['RMSprop/no_noise', 'ExtraAdam/action_noise_0.3', 'SGLD_thermal_0.01/action_noise_0.01']\n",
    "Hopper_exp = ['RMSprop/action_noise_0.2', 'ExtraAdam/action_noise_0.3', 'SGLD_thermal_0.001/action_noise_0.2']\n",
    "HalfCheetah_exp = ['RMSprop/action_noise_0.2', 'ExtraAdam/action_noise_0.01', 'SGLD_thermal_0.01/no_noise']\n",
    "Ant_exp = ['RMSprop/action_noise_0.4', 'ExtraAdam/action_noise_0.01', 'SGLD_thermal_0.0001/action_noise_0.2']\n",
    "Swimmer_exp = ['RMSprop/action_noise_0.4', 'ExtraAdam/action_noise_0.4', 'SGLD_thermal_1e-05/action_noise_0.4']\n",
    "Reacher_exp = ['RMSprop/action_noise_0.4', 'ExtraAdam/action_noise_0.2', 'SGLD_thermal_0.001/action_noise_0.2']\n",
    "Humanoid_exp = ['RMSprop/no_noise', 'ExtraAdam/action_noise_0.01', 'SGLD_thermal_0.0001/action_noise_0.01']\n",
    "InvertedPendulum_exp = ['RMSprop/action_noise_0.1', 'ExtraAdam/action_noise_0.01', 'SGLD_thermal_0.001/action_noise_0.01']\n",
    "exp = [Walker_exp, HalfCheetah_exp, Hopper_exp, Ant_exp, Swimmer_exp, Reacher_exp, Humanoid_exp, InvertedPendulum_exp]    \n",
    "plot_learning_all(exp, 'Comparison', comparison = True, OnePlayer = False, best = False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 619916 characters]\n",
      "text/plain": [
       "<Figure size 960x1296 with 8 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 496964 characters]\n",
      "text/plain": [
       "<Figure size 960x1296 with 8 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "## Comparison One Player Average\n",
    "Walker_exp = ['RMSprop/action_noise_0.01', 'ExtraAdam/action_noise_0.2', 'SGLD_thermal_0.01/action_noise_0.1']\n",
    "Hopper_exp = ['RMSprop/action_noise_0.4', 'ExtraAdam/action_noise_0.1', 'SGLD_thermal_1e-05/action_noise_0.3']\n",
    "HalfCheetah_exp = ['RMSprop/action_noise_0.4', 'ExtraAdam/action_noise_0.01', 'SGLD_thermal_0.01/action_noise_0.01']\n",
    "Ant_exp = ['RMSprop/action_noise_0.4', 'ExtraAdam/action_noise_0.01', 'SGLD_thermal_0.01/action_noise_0.4']\n",
    "Swimmer_exp = ['RMSprop/action_noise_0.3', 'ExtraAdam/action_noise_0.3', 'SGLD_thermal_0.01/action_noise_0.2']\n",
    "Reacher_exp = ['RMSprop/action_noise_0.3', 'ExtraAdam/action_noise_0.2', 'SGLD_thermal_0.001/action_noise_0.2']\n",
    "Humanoid_exp = ['RMSprop/no_noise', 'ExtraAdam/action_noise_0.01', 'SGLD_thermal_0.01/action_noise_0.1']\n",
    "InvertedPendulum_exp = ['RMSprop/action_noise_0.01', 'ExtraAdam/action_noise_0.01', 'SGLD_thermal_0.001/no_noise']\n",
    "exp = [Walker_exp, HalfCheetah_exp, Hopper_exp, Ant_exp, Swimmer_exp, Reacher_exp, Humanoid_exp, InvertedPendulum_exp] \n",
    "plot_learning_all(exp, 'Comparison', comparison = True, OnePlayer = True, best = False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "ename": "NameError",
     "evalue": "name 'plot_learning_all' is not defined",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[1;31mNameError\u001b[0m                                 Traceback (most recent call last)",
      "\u001b[1;32m<ipython-input-3-ff707a5679fd>\u001b[0m in \u001b[0;36m<module>\u001b[1;34m()\u001b[0m\n\u001b[0;32m     20\u001b[0m \u001b[0mseed\u001b[0m \u001b[1;33m=\u001b[0m \u001b[1;33m[\u001b[0m\u001b[0mWalker_seed\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mHalfCheetah_seed\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mHopper_seed\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mAnt_seed\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mSwimmer_seed\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mReacher_seed\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mHumanoid_seed\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mInvertedPendulum_seed\u001b[0m\u001b[1;33m]\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m     21\u001b[0m \u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m---> 22\u001b[1;33m \u001b[0mplot_learning_all\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mexp\u001b[0m\u001b[1;33m,\u001b[0m \u001b[1;34m'Comparison'\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mcomparison\u001b[0m \u001b[1;33m=\u001b[0m \u001b[1;32mTrue\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mOnePlayer\u001b[0m \u001b[1;33m=\u001b[0m \u001b[1;32mFalse\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mbest\u001b[0m \u001b[1;33m=\u001b[0m \u001b[1;32mTrue\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mseed\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mseed\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m",
      "\u001b[1;31mNameError\u001b[0m: name 'plot_learning_all' is not defined"
     ]
    }
   ],
   "source": [
    "## Comparison Two Player Best\n",
    "Walker_exp = ['RMSprop/no_noise', 'SGLD_thermal_0.01/action_noise_0.01']\n",
    "Hopper_exp = ['RMSprop/action_noise_0.2', 'SGLD_thermal_0.001/action_noise_0.2']\n",
    "HalfCheetah_exp = ['RMSprop/action_noise_0.2', 'SGLD_thermal_0.01/no_noise']\n",
    "Ant_exp = ['RMSprop/action_noise_0.3', 'SGLD_thermal_0.0001/action_noise_0.2']\n",
    "Swimmer_exp = ['RMSprop/no_noise', 'SGLD_thermal_1e-05/action_noise_0.4']\n",
    "Reacher_exp = ['RMSprop/action_noise_0.4', 'SGLD_thermal_0.001/action_noise_0.2']\n",
    "Humanoid_exp = ['RMSprop/no_noise', 'SGLD_thermal_0.001/no_noise']\n",
    "InvertedPendulum_exp = ['RMSprop/no_noise', 'SGLD_thermal_0.001/no_noise']\n",
    "exp = [Walker_exp, HalfCheetah_exp, Hopper_exp, Ant_exp, Swimmer_exp, Reacher_exp, Humanoid_exp, InvertedPendulum_exp]    \n",
    "\n",
    "Walker_seed = [0, 4]\n",
    "Hopper_seed = [3, 4]\n",
    "HalfCheetah_seed = [3, 2]\n",
    "Ant_seed = [1, 3]\n",
    "Swimmer_seed = [4, 2]\n",
    "Reacher_seed = [4, 2]\n",
    "Humanoid_seed = [2, 4]\n",
    "InvertedPendulum_seed = [2, 3]\n",
    "seed = [Walker_seed, HalfCheetah_seed, Hopper_seed, Ant_seed, Swimmer_seed, Reacher_seed, Humanoid_seed, InvertedPendulum_seed] \n",
    "\n",
    "plot_learning_all(exp, 'Comparison', comparison = True, OnePlayer = False, best = True, seed = seed)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "ename": "NameError",
     "evalue": "name 'plot_learning_all' is not defined",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[1;31mNameError\u001b[0m                                 Traceback (most recent call last)",
      "\u001b[1;32m<ipython-input-4-2478a0790edd>\u001b[0m in \u001b[0;36m<module>\u001b[1;34m()\u001b[0m\n\u001b[0;32m     20\u001b[0m \u001b[0mseed\u001b[0m \u001b[1;33m=\u001b[0m \u001b[1;33m[\u001b[0m\u001b[0mWalker_seed\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mHalfCheetah_seed\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mHopper_seed\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mAnt_seed\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mSwimmer_seed\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mReacher_seed\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mHumanoid_seed\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mInvertedPendulum_seed\u001b[0m\u001b[1;33m]\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m     21\u001b[0m \u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m---> 22\u001b[1;33m \u001b[0mplot_learning_all\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mexp\u001b[0m\u001b[1;33m,\u001b[0m \u001b[1;34m'Comparison'\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mcomparison\u001b[0m \u001b[1;33m=\u001b[0m \u001b[1;32mTrue\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mOnePlayer\u001b[0m \u001b[1;33m=\u001b[0m \u001b[1;32mTrue\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mbest\u001b[0m \u001b[1;33m=\u001b[0m \u001b[1;32mTrue\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mseed\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mseed\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m",
      "\u001b[1;31mNameError\u001b[0m: name 'plot_learning_all' is not defined"
     ]
    }
   ],
   "source": [
    "## Comparison One Player Best\n",
    "Walker_exp = ['RMSprop/action_noise_0.01', 'SGLD_thermal_0.01/action_noise_0.1']\n",
    "Hopper_exp = ['RMSprop/action_noise_0.3', 'SGLD_thermal_1e-05/action_noise_0.3']\n",
    "HalfCheetah_exp = ['RMSprop/action_noise_0.2', 'SGLD_thermal_0.01/no_noise']\n",
    "Ant_exp = ['RMSprop/action_noise_0.4', 'SGLD_thermal_0.01/action_noise_0.4']\n",
    "Swimmer_exp = ['RMSprop/no_noise', 'SGLD_thermal_0.0001/action_noise_0.01']\n",
    "Reacher_exp = ['RMSprop/action_noise_0.4', 'SGLD_thermal_0.001/action_noise_0.2']\n",
    "Humanoid_exp = ['RMSprop/action_noise_0.1', 'SGLD_thermal_0.01/action_noise_0.3']\n",
    "InvertedPendulum_exp = ['RMSprop/action_noise_0.01', 'SGLD_thermal_0.0001/no_noise']\n",
    "exp = [Walker_exp, HalfCheetah_exp, Hopper_exp, Ant_exp, Swimmer_exp, Reacher_exp, Humanoid_exp, InvertedPendulum_exp]    \n",
    " \n",
    "Walker_seed = [2, 0]\n",
    "Hopper_seed = [4, 2]\n",
    "HalfCheetah_seed = [4, 4]\n",
    "Ant_seed = [2, 1]\n",
    "Swimmer_seed = [2, 2]\n",
    "Reacher_seed = [2, 0]\n",
    "Humanoid_seed = [2, 3]\n",
    "InvertedPendulum_seed = [2, 3]\n",
    "seed = [Walker_seed, HalfCheetah_seed, Hopper_seed, Ant_seed, Swimmer_seed, Reacher_seed, Humanoid_seed, InvertedPendulum_seed] \n",
    "\n",
    "plot_learning_all(exp, 'Comparison', comparison = True, OnePlayer = True, best = True, seed = seed)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Mass Uncertainty Plots"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_env(env_name, params, ax, var, OnePlayer, best, seed=None):\n",
    "\n",
    "    if OnePlayer:\n",
    "        base_dir = os.getcwd() + '/models_OnePlayer/' + env_name\n",
    "    else:\n",
    "        base_dir = os.getcwd() + '/models/' + env_name\n",
    "    results = {}\n",
    "        \n",
    "    for i, exp in enumerate(params):\n",
    "        if(best):\n",
    "            upper = seed[i] + 1\n",
    "            lower = seed[i]\n",
    "        else:\n",
    "            upper = 5\n",
    "            lower = 0\n",
    "        results[exp] = {}\n",
    "        for s in range(lower, upper):\n",
    "            try:\n",
    "                if(OnePlayer):\n",
    "                    with open(base_dir + '/' + exp +'/nr_mdp_0_1/' + str(s) + '/results_model', 'rb') as f:\n",
    "                        results[exp][s] = pickle.load(f)\n",
    "                else:\n",
    "                    with open(base_dir + '/' + exp +'/nr_mdp_0.1_1/' + str(s) + '/results_model', 'rb') as f:\n",
    "                        results[exp][s] = pickle.load(f)                \n",
    "            except IOError:\n",
    "                pass\n",
    "            \n",
    "    plot_model(results, ax, env_name, var, OnePlayer)\n",
    "\n",
    "def plot_model(results, ax, env_name, var, OnePlayer):\n",
    "    x = [i for i in range(21)]\n",
    "    colors = ['#396ab1','#3e9651', '#cc2529', '#da7c30', '#94823d', '#535154', '#00FF00', '#800000', '#F08080', '#FFFF00', '#000000', '#C0C0C0']\n",
    "    \n",
    "    idx = 0\n",
    "    for exp in results:\n",
    "        if exp not in ['results_model']:\n",
    "            avg = []\n",
    "            std = []\n",
    "            for seed in results[exp]:\n",
    "                avg_seed = []\n",
    "                std_seed = []\n",
    "                for mass in results[exp][seed]:\n",
    "                    avg_seed.append(np.mean(results[exp][seed][mass]))\n",
    "                    std_seed.append(np.std(results[exp][seed][mass]) / np.sqrt(len(results[exp][seed][mass]) - 1) * 2)\n",
    "                avg.append(avg_seed)\n",
    "                std.append(std_seed)\n",
    "            avg = np.array(avg)\n",
    "            std = np.array(std)\n",
    "            avg = np.mean(avg, axis=0)\n",
    "            std = np.mean(std, axis=0)\n",
    "            \n",
    "        if 'RMSprop' in exp:\n",
    "            l = 'GAD (RMSprop)'\n",
    "        elif 'ExtraAdam' in exp:\n",
    "            l = 'Extra-Adam'\n",
    "        else:\n",
    "            l = 'MixedNE-LD (RMSProp)'                \n",
    "        if(OnePlayer):\n",
    "            ax.plot(x, avg, color=colors[idx], linestyle=\"--\")\n",
    "        else:\n",
    "            l = r'$\\beta=$'+ str(exp[45:48]) + r', $K_t$='+ str(exp[52:53])\n",
    "            ax.plot(x, avg, label=l, color=colors[idx], linestyle=\"-\")\n",
    "        if(var):\n",
    "            ax.fill_between(x, np.array(avg) - np.array(std), np.array(avg) + np.array(std), facecolor=colors[idx], alpha=0.2, interpolate=True)\n",
    "        idx += 1\n",
    "    ax.set_title(env_name)         \n",
    "    ax.set_xticks([0, 5, 10, 15, 20])\n",
    "    ax.set_xticklabels([0.5, 0.75, 1.0, 1.5, 2.0])\n",
    "    ax.set_xlabel('Relative mass')\n",
    "    ax.set_ylabel('Reward')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_env_all(params, title, comparison, OnePlayer, best, seed = None):\n",
    "    cols = 2\n",
    "    rows = 4\n",
    "    envs = ['Walker2d-v2', 'HalfCheetah-v2', 'Hopper-v2', 'Ant-v2', 'Swimmer-v2', 'Reacher-v2', 'Humanoid-v2', 'InvertedPendulum-v2']\n",
    "    fill = [True, False]\n",
    "    for f in fill:\n",
    "        fig, axs = plt.subplots(rows, cols, figsize=(20*cols/3, 9*rows/2))\n",
    "        for i, env in enumerate(envs):\n",
    "            ax = axs.flatten()[i]\n",
    "            if(comparison):\n",
    "                param = params[i]\n",
    "            else:\n",
    "                param = params\n",
    "            if(best):\n",
    "                plot_env(env, param, ax, var = f, OnePlayer = OnePlayer, best = best, seed = seed[i])\n",
    "            else:\n",
    "                plot_env(env, param, ax, var = f, OnePlayer = OnePlayer, best = best)                \n",
    "            if(comparison):\n",
    "                ax.legend(loc='upper left')\n",
    "            else:\n",
    "                if i == 0:\n",
    "                    ax.legend()\n",
    "        plt.tight_layout()\n",
    "        \n",
    "        if(OnePlayer):\n",
    "            folder = 'OnePlayer'\n",
    "        else:\n",
    "            folder = 'TwoPlayer'\n",
    "            \n",
    "        if(f):\n",
    "            if(best):\n",
    "                plt.savefig('plots/'+ folder + '/Mass_Uncertainty/' + title +'_best_fill.pdf')\n",
    "            else:\n",
    "                plt.savefig('plots/'+ folder + '/Mass_Uncertainty/' + title +'_average_fill.pdf')                    \n",
    "        else:\n",
    "            if(best):\n",
    "                plt.savefig('plots/'+ folder + '/Mass_Uncertainty/' + title +'_best.pdf')\n",
    "            else:\n",
    "                plt.savefig('plots/'+ folder + '/Mass_Uncertainty/' + title +'_average.pdf')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_env_merge(params, title, comparison):\n",
    "    cols = 2\n",
    "    rows = 4\n",
    "    envs = ['Walker2d-v2', 'HalfCheetah-v2', 'Hopper-v2', 'Ant-v2', 'Swimmer-v2', 'Reacher-v2', 'Humanoid-v2', 'InvertedPendulum-v2']\n",
    "    fill = [True, False]\n",
    "    for f in fill:\n",
    "        fig, axs = plt.subplots(rows, cols, figsize=(20*cols/3, 9*rows/2))\n",
    "        for i, env in enumerate(envs):\n",
    "            ax = axs.flatten()[i]\n",
    "            if(comparison):\n",
    "                param = params[i]\n",
    "            else:\n",
    "                param = params\n",
    "            plot_env(env, param[0:3], ax, var = f, OnePlayer = True, best = False)\n",
    "            plot_env(env, param[3:6], ax, var = f, OnePlayer = False, best = False)                \n",
    "            if(comparison):\n",
    "                ax.legend(loc='upper left')\n",
    "            else:\n",
    "                if i == 0:\n",
    "                    ax.legend()\n",
    "        plt.tight_layout()\n",
    "        if(f):\n",
    "            plt.savefig('plots/Merged/Mass_Uncertainty_fill.pdf')\n",
    "        else:\n",
    "            plt.savefig('plots/Merged/Mass_Uncertainty.pdf')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Merged Results for OnePlayer and TwoPlayer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 530356 characters]\n",
      "text/plain": [
       "<Figure size 960x1296 with 8 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 380940 characters]\n",
      "text/plain": [
       "<Figure size 960x1296 with 8 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "Walker_exp = ['RMSprop/action_noise_0.01', 'ExtraAdam/action_noise_0.2', 'SGLD_thermal_0.01/action_noise_0.1',\\\n",
    "             'RMSprop/no_noise','ExtraAdam/action_noise_0.3', 'SGLD_thermal_0.01/action_noise_0.01']\n",
    "\n",
    "Hopper_exp = ['RMSprop/action_noise_0.4', 'ExtraAdam/action_noise_0.1', 'SGLD_thermal_1e-05/action_noise_0.3',\n",
    "               'RMSprop/action_noise_0.2', 'ExtraAdam/action_noise_0.3', 'SGLD_thermal_0.001/action_noise_0.2']\n",
    "\n",
    "HalfCheetah_exp = ['RMSprop/action_noise_0.4', 'ExtraAdam/action_noise_0.01', 'SGLD_thermal_0.01/action_noise_0.01',\\\n",
    "                   'RMSprop/action_noise_0.2', 'ExtraAdam/action_noise_0.01', 'SGLD_thermal_0.01/no_noise']\n",
    "\n",
    "Ant_exp = ['RMSprop/action_noise_0.4', 'ExtraAdam/action_noise_0.01', 'SGLD_thermal_0.01/action_noise_0.4',\\\n",
    "        'RMSprop/action_noise_0.4', 'ExtraAdam/action_noise_0.01', 'SGLD_thermal_0.0001/action_noise_0.2']\n",
    "\n",
    "Swimmer_exp = ['RMSprop/action_noise_0.3', 'ExtraAdam/action_noise_0.3', 'SGLD_thermal_0.01/action_noise_0.2',\\\n",
    "               'RMSprop/action_noise_0.4', 'ExtraAdam/action_noise_0.4', 'SGLD_thermal_1e-05/action_noise_0.4']\n",
    "\n",
    "Reacher_exp = ['RMSprop/action_noise_0.3', 'ExtraAdam/action_noise_0.2', 'SGLD_thermal_0.001/action_noise_0.2',\\\n",
    "               'RMSprop/action_noise_0.4', 'ExtraAdam/action_noise_0.2', 'SGLD_thermal_0.001/action_noise_0.2']\n",
    "\n",
    "Humanoid_exp = ['RMSprop/no_noise', 'ExtraAdam/action_noise_0.01', 'SGLD_thermal_0.01/action_noise_0.1',\\\n",
    "                'RMSprop/no_noise', 'ExtraAdam/action_noise_0.01', 'SGLD_thermal_0.0001/action_noise_0.01']\n",
    "\n",
    "InvertedPendulum_exp = ['RMSprop/action_noise_0.01', 'ExtraAdam/action_noise_0.01', 'SGLD_thermal_0.001/no_noise',\\\n",
    "                        'RMSprop/action_noise_0.1', 'ExtraAdam/action_noise_0.01', 'SGLD_thermal_0.001/action_noise_0.01']\n",
    "params = [Walker_exp, HalfCheetah_exp, Hopper_exp, Ant_exp, Swimmer_exp, Reacher_exp, Humanoid_exp, InvertedPendulum_exp]    \n",
    "plot_env_merge(params, title = 'Comparison', comparison = True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Average Results for OnePlayer and TwoPlayer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 358792 characters]\n",
      "text/plain": [
       "<Figure size 960x1296 with 8 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 255584 characters]\n",
      "text/plain": [
       "<Figure size 960x1296 with 8 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "## Comparison Two Player Average\n",
    "Walker_exp = ['RMSprop/no_noise','ExtraAdam/action_noise_0.3', 'SGLD_thermal_0.01/action_noise_0.01']\n",
    "Hopper_exp = ['RMSprop/action_noise_0.2', 'ExtraAdam/action_noise_0.3', 'SGLD_thermal_0.001/action_noise_0.2']\n",
    "HalfCheetah_exp = ['RMSprop/action_noise_0.2', 'ExtraAdam/action_noise_0.01', 'SGLD_thermal_0.01/no_noise']\n",
    "Ant_exp = ['RMSprop/action_noise_0.4', 'ExtraAdam/action_noise_0.01', 'SGLD_thermal_0.0001/action_noise_0.2']\n",
    "Swimmer_exp = ['RMSprop/action_noise_0.4', 'ExtraAdam/action_noise_0.4', 'SGLD_thermal_1e-05/action_noise_0.4']\n",
    "Reacher_exp = ['RMSprop/action_noise_0.4', 'ExtraAdam/action_noise_0.2', 'SGLD_thermal_0.001/action_noise_0.2']\n",
    "Humanoid_exp = ['RMSprop/no_noise', 'ExtraAdam/action_noise_0.01', 'SGLD_thermal_0.0001/action_noise_0.01']\n",
    "InvertedPendulum_exp = ['RMSprop/action_noise_0.1', 'ExtraAdam/action_noise_0.01', 'SGLD_thermal_0.001/action_noise_0.01']\n",
    "params = [Walker_exp, HalfCheetah_exp, Hopper_exp, Ant_exp, Swimmer_exp, Reacher_exp, Humanoid_exp, InvertedPendulum_exp]    \n",
    "plot_env_all(params, title = 'Comparison', comparison = True, OnePlayer = False, best = False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 374840 characters]\n",
      "text/plain": [
       "<Figure size 960x1296 with 8 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 277968 characters]\n",
      "text/plain": [
       "<Figure size 960x1296 with 8 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "## Comparison One Player Average\n",
    "Walker_exp = ['RMSprop/action_noise_0.01', 'ExtraAdam/action_noise_0.2', 'SGLD_thermal_0.01/action_noise_0.1']\n",
    "Hopper_exp = ['RMSprop/action_noise_0.4', 'ExtraAdam/action_noise_0.1', 'SGLD_thermal_1e-05/action_noise_0.3']\n",
    "HalfCheetah_exp = ['RMSprop/action_noise_0.4', 'ExtraAdam/action_noise_0.01', 'SGLD_thermal_0.01/action_noise_0.01']\n",
    "Ant_exp = ['RMSprop/action_noise_0.4', 'ExtraAdam/action_noise_0.01', 'SGLD_thermal_0.01/action_noise_0.4']\n",
    "Swimmer_exp = ['RMSprop/action_noise_0.3', 'ExtraAdam/action_noise_0.3', 'SGLD_thermal_0.01/action_noise_0.2']\n",
    "Reacher_exp = ['RMSprop/action_noise_0.3', 'ExtraAdam/action_noise_0.2', 'SGLD_thermal_0.001/action_noise_0.2']\n",
    "Humanoid_exp = ['RMSprop/no_noise', 'ExtraAdam/action_noise_0.01', 'SGLD_thermal_0.01/action_noise_0.1']\n",
    "InvertedPendulum_exp = ['RMSprop/action_noise_0.01', 'ExtraAdam/action_noise_0.01', 'SGLD_thermal_0.001/no_noise']\n",
    "params = [Walker_exp, HalfCheetah_exp, Hopper_exp, Ant_exp, Swimmer_exp, Reacher_exp, Humanoid_exp, InvertedPendulum_exp]    \n",
    "plot_env_all(params, 'Comparison', comparison = True, OnePlayer = True, best = False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Best Seed Results for OnePlayer and TwoPlayer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "## Comparison Two Player Best\n",
    "Walker_exp = ['RMSprop/no_noise', 'SGLD_thermal_0.01/action_noise_0.01']\n",
    "Hopper_exp = ['RMSprop/action_noise_0.2', 'SGLD_thermal_0.001/action_noise_0.2']\n",
    "HalfCheetah_exp = ['RMSprop/action_noise_0.2', 'SGLD_thermal_0.01/no_noise']\n",
    "Ant_exp = ['RMSprop/action_noise_0.3', 'SGLD_thermal_0.0001/action_noise_0.2']\n",
    "Swimmer_exp = ['RMSprop/no_noise', 'SGLD_thermal_1e-05/action_noise_0.4']\n",
    "Reacher_exp = ['RMSprop/action_noise_0.4', 'SGLD_thermal_0.001/action_noise_0.2']\n",
    "Humanoid_exp = ['RMSprop/no_noise', 'SGLD_thermal_0.001/no_noise']\n",
    "InvertedPendulum_exp = ['RMSprop/no_noise', 'SGLD_thermal_0.001/no_noise']\n",
    "exp = [Walker_exp, HalfCheetah_exp, Hopper_exp, Ant_exp, Swimmer_exp, Reacher_exp, Humanoid_exp, InvertedPendulum_exp]    \n",
    "Walker_seed = [0, 4]\n",
    "Hopper_seed = [3, 4]\n",
    "HalfCheetah_seed = [3, 2]\n",
    "Ant_seed = [1, 3]\n",
    "Swimmer_seed = [4, 2]\n",
    "Reacher_seed = [4, 2]\n",
    "Humanoid_seed = [2, 4]\n",
    "InvertedPendulum_seed = [2, 3]\n",
    "\n",
    "seed = [Walker_seed, HalfCheetah_seed, Hopper_seed, Ant_seed, Swimmer_seed, Reacher_seed, Humanoid_seed, InvertedPendulum_seed] \n",
    "plot_env_all(exp, 'Comparison', comparison = True, OnePlayer = False, best = True, seed = seed)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/yuyu/anaconda3/lib/python3.6/site-packages/matplotlib/pyplot.py:537: RuntimeWarning: More than 20 figures have been opened. Figures created through the pyplot interface (`matplotlib.pyplot.figure`) are retained until explicitly closed and may consume too much memory. (To control this warning, see the rcParam `figure.max_open_warning`).\n",
      "  max_open_warning, RuntimeWarning)\n"
     ]
    }
   ],
   "source": [
    "## Comparison One Player Best\n",
    "Walker_exp = ['RMSprop/action_noise_0.01', 'SGLD_thermal_0.01/action_noise_0.1']\n",
    "Hopper_exp = ['RMSprop/action_noise_0.3', 'SGLD_thermal_1e-05/action_noise_0.3']\n",
    "HalfCheetah_exp = ['RMSprop/action_noise_0.2', 'SGLD_thermal_0.01/no_noise']\n",
    "Ant_exp = ['RMSprop/action_noise_0.4', 'SGLD_thermal_0.01/action_noise_0.4']\n",
    "Swimmer_exp = ['RMSprop/no_noise', 'SGLD_thermal_0.0001/action_noise_0.01']\n",
    "Reacher_exp = ['RMSprop/action_noise_0.4', 'SGLD_thermal_0.001/action_noise_0.2']\n",
    "Humanoid_exp = ['RMSprop/action_noise_0.1', 'SGLD_thermal_0.01/action_noise_0.3']\n",
    "InvertedPendulum_exp = ['RMSprop/action_noise_0.01', 'SGLD_thermal_0.0001/no_noise']\n",
    "exp = [Walker_exp, HalfCheetah_exp, Hopper_exp, Ant_exp, Swimmer_exp, Reacher_exp, Humanoid_exp, InvertedPendulum_exp]    \n",
    " \n",
    "Walker_seed = [2, 0]\n",
    "Hopper_seed = [4, 2]\n",
    "HalfCheetah_seed = [4, 4]\n",
    "Ant_seed = [2, 1]\n",
    "Swimmer_seed = [2, 2]\n",
    "Reacher_seed = [2, 0]\n",
    "Humanoid_seed = [2, 3]\n",
    "InvertedPendulum_seed = [2, 3]\n",
    "seed = [Walker_seed, HalfCheetah_seed, Hopper_seed, Ant_seed, Swimmer_seed, Reacher_seed, Humanoid_seed, InvertedPendulum_seed] \n",
    "\n",
    "plot_env_all(exp, 'Comparison', comparison = True, OnePlayer = True, best = True, seed = seed)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Ablation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_ablation(env_name, params, ax, var):\n",
    "\n",
    "    base_dir = os.getcwd() + '/models/' + env_name\n",
    "    results = {}\n",
    "        \n",
    "    for i, exp in enumerate(params):\n",
    "        results[exp] = {}\n",
    "        for s in range(5):\n",
    "            try:\n",
    "                with open(base_dir + '/' + exp + str(s) + '/results_model', 'rb') as f:\n",
    "                    results[exp][s] = pickle.load(f)             \n",
    "            except IOError:\n",
    "                pass\n",
    "            \n",
    "    plot_model(results, ax, env_name, var, OnePlayer=False)\n",
    "\n",
    "def plot_ablation_all(params, title, comparison):\n",
    "    cols = 1\n",
    "    rows = 1\n",
    "    envs = ['HalfCheetah-v2']\n",
    "    fill = [True, False]\n",
    "    for f in fill:\n",
    "        fig, axs = plt.subplots(rows, cols, figsize=(20*cols/3, 9*rows/2))\n",
    "        for i, env in enumerate(envs):\n",
    "            ax = axs\n",
    "            if(comparison):\n",
    "                param = params[i]\n",
    "            else:\n",
    "                param = params\n",
    "            plot_ablation(env, param, ax, var = f)                \n",
    "            if(comparison):\n",
    "                ax.legend(loc='upper left')\n",
    "            else:\n",
    "                if i == 0:\n",
    "                    ax.legend()\n",
    "        plt.tight_layout()\n",
    "        if(f):\n",
    "            plt.savefig('plots/TwoPlayer/Mass_Uncertainty/Ablation_fill.pdf')\n",
    "        else:\n",
    "            plt.savefig('plots/TwoPlayer/Mass_Uncertainty/Ablation.pdf')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 95652 characters]\n",
      "text/plain": [
       "<Figure size 480x324 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 71452 characters]\n",
      "text/plain": [
       "<Figure size 480x324 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "param = 'SGLD_thermal_0.01/no_noise/nr_mdp_0.1_1/'\n",
    "_exp = [param+'beta_0.5/Kt_1/', param+'beta_0.5/Kt_2/', param+'beta_0.5/Kt_5/', param+'beta_0.9/Kt_1/', param+'beta_0.9/Kt_2/',\\\n",
    "              param+'beta_0.9/Kt_5/', param+'beta_1.0/Kt_1/', param+'beta_1.0/Kt_2/', param+'beta_1.0/Kt_5/']\n",
    "exp = [_exp]    \n",
    "plot_ablation_all(exp, title = 'Ablation', comparison = True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": true
   },
   "source": [
    "## Heatmap"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "from mpl_toolkits.axes_grid1 import make_axes_locatable\n",
    "\n",
    "def plot_2d(env_name, params, OnePlayer, best, seed=None):\n",
    "    if OnePlayer:\n",
    "        base_dir = os.getcwd() + '/models_OnePlayer/' + env_name\n",
    "    else:\n",
    "        base_dir = os.getcwd() + '/models/' + env_name\n",
    "\n",
    "    results = {}\n",
    "    vmin = 2000\n",
    "    vmax = -2000\n",
    "    for idx, param in enumerate(params):\n",
    "        if(best):\n",
    "            upper = seed[idx] + 1\n",
    "            lower = seed[idx]\n",
    "        else:\n",
    "            upper = 5\n",
    "            lower = 0\n",
    "        try:\n",
    "            mat = np.zeros((11, 11))\n",
    "            for s in range(lower, upper):\n",
    "                if(OnePlayer):\n",
    "                    with open(base_dir + '/' + param + '/nr_mdp_0_1/'+str(s)+'/results_model_noise', 'rb') as f:\n",
    "                        data = pickle.load(f)\n",
    "                else:\n",
    "                    with open(base_dir + '/' + param + '/nr_mdp_0.1_1/'+str(s)+'/results_model_noise', 'rb') as f:\n",
    "                        data = pickle.load(f)                  \n",
    "                row = 0\n",
    "                for key in data:\n",
    "                    col = 0\n",
    "                    for key2 in data[key]:\n",
    "                        mat[row, col] += np.mean(data[key][key2])\n",
    "                        col += 1\n",
    "                    row += 1\n",
    "            mat = mat/5\n",
    "            results[param] = mat\n",
    "\n",
    "            vmin = min(vmin, np.min(mat))\n",
    "            vmax = max(vmax, np.max(mat))\n",
    "        except:\n",
    "            pass\n",
    "    \n",
    "    plt.figure(figsize=(7, 7))\n",
    "    idx = 1\n",
    "    for key in results:\n",
    "        \n",
    "        plt.subplot(1, 4, idx)\n",
    "        idx += 1\n",
    "        ax = plt.gca()\n",
    "        im = ax.imshow(results[key], cmap='autumn', interpolation='nearest', vmin=vmin, vmax=vmax)\n",
    "        x = [0,2,4,6,8,10]\n",
    "        #[0, 0.05, 0.1, 0.15, 0.2, 0.25, 0.3, 0.35, 0.4, 0.45, 0.5]\n",
    "        labels = [0, 0.1, 0.2, 0.3, 0.4, 0.5]\n",
    "        plt.xticks(x, labels, fontsize=4)\n",
    "\n",
    "        y = [0,2,4,6,8,10]\n",
    "        #[0.5, 0.6, 0.7, 0.8, 0.9, 1.0, 1.2, 1.4, 1.6, 1.8, 2.0]\n",
    "        labels = [0.5, 0.7, 0.9, 1.2, 1.6, 2.0]\n",
    "        plt.yticks(y, labels, fontsize=4)\n",
    "\n",
    "        plt.xlabel('Noise Probability', fontsize=5)\n",
    "        plt.ylabel('Relative Mass', fontsize=5)\n",
    "        if \"RMSprop\" in key:\n",
    "            plt.title('GAD' + '/' + env_name, fontsize=6)\n",
    "        elif \"ExtraAdam\" in key:\n",
    "            plt.title('Extra-Adam' + '/' + env_name, fontsize=6)\n",
    "        else:\n",
    "            plt.title('MixedNE-LD' + '/' + env_name, fontsize=6)            \n",
    "        divider = make_axes_locatable(ax)\n",
    "        cax = divider.append_axes(\"right\", size=\"5%\", pad=0.13)\n",
    "\n",
    "        cbar = plt.colorbar(im, cax=cax)\n",
    "        cbar.ax.tick_params(labelsize=4) \n",
    "    plt.tight_layout()\n",
    "    \n",
    "    if(OnePlayer):\n",
    "        folder = 'OnePlayer'\n",
    "    else:\n",
    "        folder = 'TwoPlayer'\n",
    "    \n",
    "    if(best):\n",
    "        plt.savefig('plots/'+ folder + '/Heat_map/' + env_name +'_best.pdf', bbox_inches='tight',pad_inches=0.0, dpi= 300)\n",
    "    else:\n",
    "        plt.savefig('plots/'+ folder + '/Heat_map/' + env_name +'_average.pdf', bbox_inches='tight',pad_inches=0.0, dpi = 300) "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Average results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## Comparison Two Player Average\n",
    "Walker_exp = ['RMSprop/no_noise','ExtraAdam/action_noise_0.3', 'SGLD_thermal_0.01/action_noise_0.01']\n",
    "Hopper_exp = ['RMSprop/action_noise_0.2', 'ExtraAdam/action_noise_0.3', 'SGLD_thermal_0.001/action_noise_0.2']\n",
    "HalfCheetah_exp = ['RMSprop/action_noise_0.2', 'ExtraAdam/action_noise_0.01', 'SGLD_thermal_0.01/no_noise']\n",
    "Ant_exp = ['RMSprop/action_noise_0.4', 'ExtraAdam/action_noise_0.01', 'SGLD_thermal_0.0001/action_noise_0.2']\n",
    "Swimmer_exp = ['RMSprop/action_noise_0.4', 'ExtraAdam/action_noise_0.4', 'SGLD_thermal_1e-05/action_noise_0.4']\n",
    "Reacher_exp = ['RMSprop/action_noise_0.4', 'ExtraAdam/action_noise_0.2', 'SGLD_thermal_0.001/action_noise_0.2']\n",
    "Humanoid_exp = ['RMSprop/no_noise', 'ExtraAdam/action_noise_0.01', 'SGLD_thermal_0.0001/action_noise_0.01']\n",
    "InvertedPendulum_exp = ['RMSprop/action_noise_0.1', 'ExtraAdam/action_noise_0.01', 'SGLD_thermal_0.001/action_noise_0.01']\n",
    "\n",
    "plot_2d('Walker2d-v2', Walker_exp, OnePlayer = False, best = False)\n",
    "plot_2d('Hopper-v2', Hopper_exp, OnePlayer = False, best = False)\n",
    "plot_2d('HalfCheetah-v2', HalfCheetah_exp, OnePlayer = False, best = False)\n",
    "plot_2d('Ant-v2', Ant_exp, OnePlayer = False, best = False)\n",
    "plot_2d('Swimmer-v2', Swimmer_exp, OnePlayer = False, best = False)\n",
    "plot_2d('Reacher-v2', Reacher_exp, OnePlayer = False, best = False)\n",
    "plot_2d('Humanoid-v2', Humanoid_exp, OnePlayer = False, best = False)\n",
    "plot_2d('InvertedPendulum-v2', InvertedPendulum_exp, OnePlayer = False, best = False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## Comparison One Player Average\n",
    "Walker_exp = ['RMSprop/action_noise_0.01', 'ExtraAdam/action_noise_0.2', 'SGLD_thermal_0.01/action_noise_0.1']\n",
    "Hopper_exp = ['RMSprop/action_noise_0.4', 'ExtraAdam/action_noise_0.1', 'SGLD_thermal_1e-05/action_noise_0.3']\n",
    "HalfCheetah_exp = ['RMSprop/action_noise_0.4', 'ExtraAdam/action_noise_0.01', 'SGLD_thermal_0.01/action_noise_0.01']\n",
    "Ant_exp = ['RMSprop/action_noise_0.4', 'ExtraAdam/action_noise_0.01', 'SGLD_thermal_0.01/action_noise_0.4']\n",
    "Swimmer_exp = ['RMSprop/action_noise_0.3', 'ExtraAdam/action_noise_0.3', 'SGLD_thermal_0.01/action_noise_0.2']\n",
    "Reacher_exp = ['RMSprop/action_noise_0.3', 'ExtraAdam/action_noise_0.2', 'SGLD_thermal_0.001/action_noise_0.2']\n",
    "Humanoid_exp = ['RMSprop/no_noise', 'ExtraAdam/action_noise_0.01', 'SGLD_thermal_0.01/action_noise_0.1']\n",
    "InvertedPendulum_exp = ['RMSprop/action_noise_0.01', 'ExtraAdam/action_noise_0.01', 'SGLD_thermal_0.001/no_noise']\n",
    "\n",
    "plot_2d('Walker2d-v2', Walker_exp, OnePlayer = True, best = False)\n",
    "plot_2d('Hopper-v2', Hopper_exp, OnePlayer = True, best = False)\n",
    "plot_2d('HalfCheetah-v2', HalfCheetah_exp, OnePlayer = True, best = False)\n",
    "plot_2d('Ant-v2', Ant_exp, OnePlayer = True, best = False)\n",
    "plot_2d('Swimmer-v2', Swimmer_exp, OnePlayer = True, best = False)\n",
    "plot_2d('Reacher-v2', Reacher_exp, OnePlayer = True, best = False)\n",
    "plot_2d('Humanoid-v2', Humanoid_exp, OnePlayer = True, best = False)\n",
    "plot_2d('InvertedPendulum-v2', InvertedPendulum_exp, OnePlayer = True, best = False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Best seed results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## Comparison Two Player Best\n",
    "Walker_param = ['RMSprop/no_noise', 'SGLD_thermal_0.01/action_noise_0.01']\n",
    "Hopper_param = ['RMSprop/action_noise_0.2', 'SGLD_thermal_0.001/action_noise_0.2']\n",
    "HalfCheetah_param = ['RMSprop/action_noise_0.2', 'SGLD_thermal_0.01/no_noise']\n",
    "Ant_param = ['RMSprop/action_noise_0.3', 'SGLD_thermal_0.0001/action_noise_0.2']\n",
    "Swimmer_param = ['RMSprop/no_noise', 'SGLD_thermal_1e-05/action_noise_0.4']\n",
    "Reacher_param = ['RMSprop/action_noise_0.4', 'SGLD_thermal_0.001/action_noise_0.2']\n",
    "Humanoid_param = ['RMSprop/no_noise', 'SGLD_thermal_0.001/no_noise']\n",
    "InvertedPendulum_param = ['RMSprop/no_noise', 'SGLD_thermal_0.001/no_noise']\n",
    "\n",
    "Walker_seed = [0, 4]\n",
    "Hopper_seed = [3, 4]\n",
    "HalfCheetah_seed = [3, 2]\n",
    "Ant_seed = [1, 3]\n",
    "Swimmer_seed = [4, 2]\n",
    "Reacher_seed = [4, 2]\n",
    "Humanoid_seed = [2, 4]\n",
    "InvertedPendulum_seed = [2, 3]\n",
    "\n",
    "\n",
    "plot_2d('Walker2d-v2', Walker_param, OnePlayer = False, best = True, seed = Walker_seed)\n",
    "plot_2d('Hopper-v2', Hopper_param, OnePlayer = False, best = True, seed = Hopper_seed)\n",
    "plot_2d('HalfCheetah-v2', HalfCheetah_param, OnePlayer = False, best = True, seed = HalfCheetah_seed)\n",
    "plot_2d('Ant-v2', Ant_param, OnePlayer = False, best = True, seed = Ant_seed)\n",
    "plot_2d('Swimmer-v2', Swimmer_param, OnePlayer = False, best = True, seed = Swimmer_seed)\n",
    "plot_2d('Reacher-v2', Reacher_param, OnePlayer = False, best = True, seed = Reacher_seed)\n",
    "plot_2d('Humanoid-v2', Humanoid_param, OnePlayer = False, best = True, seed = Humanoid_seed)\n",
    "plot_2d('InvertedPendulum-v2', InvertedPendulum_param, OnePlayer = False, best = True, seed = InvertedPendulum_seed)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## Comparison One Player Best\n",
    "Walker_param = ['RMSprop/action_noise_0.01', 'SGLD_thermal_0.01/action_noise_0.1']\n",
    "Hopper_param = ['RMSprop/action_noise_0.3', 'SGLD_thermal_1e-05/action_noise_0.3']\n",
    "HalfCheetah_param = ['RMSprop/action_noise_0.2', 'SGLD_thermal_0.01/no_noise']\n",
    "Ant_param = ['RMSprop/action_noise_0.4', 'SGLD_thermal_0.01/action_noise_0.4']\n",
    "Swimmer_param = ['RMSprop/no_noise', 'SGLD_thermal_0.0001/action_noise_0.01']\n",
    "Reacher_param = ['RMSprop/action_noise_0.4', 'SGLD_thermal_0.001/action_noise_0.2']\n",
    "Humanoid_param = ['RMSprop/action_noise_0.1', 'SGLD_thermal_0.01/action_noise_0.3']\n",
    "InvertedPendulum_param = ['RMSprop/action_noise_0.01', 'SGLD_thermal_0.0001/no_noise']\n",
    "\n",
    "Walker_seed = [2, 0]\n",
    "Hopper_seed = [4, 2]\n",
    "HalfCheetah_seed = [4, 4]\n",
    "Ant_seed = [2, 1]\n",
    "Swimmer_seed = [2, 2]\n",
    "Reacher_seed = [2, 0]\n",
    "Humanoid_seed = [2, 3]\n",
    "InvertedPendulum_seed = [2, 3]\n",
    "\n",
    "plot_2d('Walker2d-v2', Walker_param, OnePlayer = True, best = True, seed = Walker_seed)\n",
    "plot_2d('Hopper-v2', Hopper_param, OnePlayer = True, best = True, seed = Hopper_seed)\n",
    "plot_2d('HalfCheetah-v2', HalfCheetah_param, OnePlayer = True, best = True, seed = HalfCheetah_seed)\n",
    "plot_2d('Ant-v2', Ant_param, OnePlayer = True, best = True, seed = Ant_seed)\n",
    "plot_2d('Swimmer-v2', Swimmer_param, OnePlayer = True, best = True, seed = Swimmer_seed)\n",
    "plot_2d('Reacher-v2', Reacher_param, OnePlayer = True, best = True, seed = Reacher_seed)\n",
    "plot_2d('Humanoid-v2', Humanoid_param, OnePlayer = True, best = True, seed = Humanoid_seed)\n",
    "plot_2d('InvertedPendulum-v2', InvertedPendulum_param, OnePlayer = True, best = True, seed = InvertedPendulum_seed)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Ablation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 58,
   "metadata": {},
   "outputs": [],
   "source": [
    "from mpl_toolkits.axes_grid1 import make_axes_locatable\n",
    "\n",
    "def plot_2d(env_name, params):\n",
    "\n",
    "    base_dir = os.getcwd() + '/models/' + env_name\n",
    "\n",
    "    results = {}\n",
    "    vmin = 2000\n",
    "    vmax = -2000\n",
    "    for idx, param in enumerate(params):\n",
    "        try:\n",
    "            mat = np.zeros((11, 11))\n",
    "            for s in range(5):\n",
    "                with open(base_dir + '/' + param +str(s)+'/results_model_noise', 'rb') as f:\n",
    "                    data = pickle.load(f)                  \n",
    "                row = 0\n",
    "                for key in data:\n",
    "                    col = 0\n",
    "                    for key2 in data[key]:\n",
    "                        mat[row, col] += np.mean(data[key][key2])\n",
    "                        col += 1\n",
    "                    row += 1\n",
    "            mat = mat/5\n",
    "            results[param] = mat\n",
    "\n",
    "            vmin = min(vmin, np.min(mat))\n",
    "            vmax = max(vmax, np.max(mat))\n",
    "        except:\n",
    "            pass\n",
    "    \n",
    "    plt.figure(figsize=(8, 8))\n",
    "    idx = 1\n",
    "    for exp in results:\n",
    "        \n",
    "        plt.subplot(3, 3, idx)\n",
    "        idx += 1\n",
    "        ax = plt.gca()\n",
    "        im = ax.imshow(results[exp], cmap='autumn', interpolation='nearest', vmin=vmin, vmax=vmax)\n",
    "        x = [0,2,4,6,8,10]\n",
    "        #[0, 0.05, 0.1, 0.15, 0.2, 0.25, 0.3, 0.35, 0.4, 0.45, 0.5]\n",
    "        labels = [0, 0.1, 0.2, 0.3, 0.4, 0.5]\n",
    "        plt.xticks(x, labels, fontsize=6)\n",
    "\n",
    "        y = [0,2,4,6,8,10]\n",
    "        #[0.5, 0.6, 0.7, 0.8, 0.9, 1.0, 1.2, 1.4, 1.6, 1.8, 2.0]\n",
    "        labels = [0.5, 0.7, 0.9, 1.2, 1.6, 2.0]\n",
    "        plt.yticks(y, labels, fontsize=6)\n",
    "\n",
    "        plt.xlabel('Noise Probability', fontsize=7)\n",
    "        plt.ylabel('Relative Mass', fontsize=7)\n",
    "        l = r'$\\beta=$'+ str(exp[45:48]) + r', $K_t$='+ str(exp[52:53])\n",
    "        plt.title(l, fontsize=10)            \n",
    "        divider = make_axes_locatable(ax)\n",
    "        cax = divider.append_axes(\"right\", size=\"5%\", pad=0.13)\n",
    "\n",
    "        cbar = plt.colorbar(im, cax=cax)\n",
    "        cbar.ax.tick_params(labelsize=6) \n",
    "    plt.tight_layout()\n",
    "    plt.savefig('plots/TwoPlayer/Heat_map/Ablation.pdf')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 59,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 30900 characters]\n",
      "text/plain": [
       "<Figure size 576x576 with 18 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "param = 'SGLD_thermal_0.01/no_noise/nr_mdp_0.1_1/'\n",
    "exp = [param+'beta_0.5/Kt_1/', param+'beta_0.5/Kt_2/', param+'beta_0.5/Kt_5/', param+'beta_0.9/Kt_1/', param+'beta_0.9/Kt_2/',\\\n",
    "              param+'beta_0.9/Kt_5/', param+'beta_1.0/Kt_1/', param+'beta_1.0/Kt_2/', param+'beta_1.0/Kt_5/']\n",
    "plot_2d('HalfCheetah-v2', exp)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
