{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import os\n",
    "import pickle"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Learning Curve"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 93,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_learning(env_name, ax, experiment, var = True, OnePlayer = False):\n",
    "    if OnePlayer:\n",
    "        base_dir = os.getcwd() + '/models_OnePlayer/' + env_name\n",
    "    else:\n",
    "        base_dir = os.getcwd() + '/models_TwoPlayer/' + env_name\n",
    "    results = {}\n",
    "        \n",
    "    for exp in experiment:\n",
    "        results[exp] = {}\n",
    "        for seed in range(1):\n",
    "            try:\n",
    "                if(OnePlayer):\n",
    "                    results[exp][seed] = np.load(base_dir + '/' + exp + '/alpha_0.0/' + str(seed) + '/results.npy')\n",
    "                else:\n",
    "                    results[exp][seed] = np.load(base_dir + '/' + exp + str(seed) + '/results.npy')   \n",
    "            except IOError:\n",
    "                print(\"No Files found\")\n",
    "            \n",
    "    plot_learning_curves(results, ax, env_name, var, OnePlayer)\n",
    "\n",
    "def plot_learning_curves(results, ax, env_name, var, OnePlayer):    \n",
    "    colors = ['#396ab1', '#3e9651', '#cc2529', '#da7c30', '#94823d', '#535154', '#006400', '#00FF00', '#800000', '#F08080', '#FFFF00', '#000000', '#C0C0C0']\n",
    "    idx = 0\n",
    "    for exp in results:\n",
    "        for seed in results[exp]:\n",
    "            data = results[exp][seed]\n",
    "            x = np.linspace(0, 5e3*(len(data)-1), len(data))\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, data.tolist(), color = colors[idx], label = exp)\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))    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 94,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_learning_all(experiment, title, comparison, OnePlayer):\n",
    "    cols = 3\n",
    "    rows = 4 \n",
    "    envs = ['Walker2d-v2', 'Walker2d-v2', 'Walker2d-v2', \\\n",
    "            'HalfCheetah-v2', 'HalfCheetah-v2', 'HalfCheetah-v2', \\\n",
    "            'Humanoid-v2', 'Humanoid-v2', 'Humanoid-v2',\\\n",
    "            'InvertedPendulum-v2', 'InvertedPendulum-v2', 'InvertedPendulum-v2'\n",
    "            ]\n",
    "    fill = [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",
    "\n",
    "            if(i%3 == 0):\n",
    "                plot_learning(env, ax, exp, var = f, OnePlayer = True)\n",
    "            else:\n",
    "                plot_learning(env, ax, exp, var = f, OnePlayer = False)\n",
    "                    \n",
    "            if(comparison):\n",
    "                ax.legend()\n",
    "            else:\n",
    "                if i == 0:\n",
    "                    ax.legend()\n",
    "        plt.tight_layout()\n",
    "#         plt.savefig('plots/Learning_curve.pdf')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 95,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 655032 characters]\n",
      "text/plain": [
       "<Figure size 1440x1296 with 12 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "Walker_0 = ['RMSprop/action_noise_0.01', 'ExtraAdam/action_noise_0.1', 'SGLD_thermal_1e-05/action_noise_0.01']\n",
    "Walker_1 = ['RMSprop/action_noise_0.01/alpha_0.1/', 'ExtraAdam/action_noise_0.01/alpha_0.1/','SGLD_thermal_0.001/action_noise_0.01/alpha_0.1/']\n",
    "Walker_2 = ['RMSprop/action_noise_0.005/alpha_0.2/', 'ExtraAdam/action_noise_0.005/alpha_0.2/','SGLD_thermal_0.0001/action_noise_0.005/alpha_0.2/']\n",
    "\n",
    "HalfCheetah_0 = ['RMSprop/action_noise_0.01', 'ExtraAdam/action_noise_0.001', 'SGLD_thermal_1e-05/action_noise_0.01']\n",
    "HalfCheetah_1 = ['RMSprop/action_noise_0.01/alpha_0.1/', 'ExtraAdam/action_noise_0.01/alpha_0.1/', 'SGLD_thermal_1e-05/action_noise_0.1/alpha_0.1/']\n",
    "HalfCheetah_2 = ['RMSprop/action_noise_0.005/alpha_0.2/', 'ExtraAdam/action_noise_0.1/alpha_0.2/', 'SGLD_thermal_1e-05/action_noise_0.005/alpha_0.2/']  \n",
    "\n",
    "Humanoid_0 = ['RMSprop/action_noise_0.1', 'ExtraAdam/action_noise_0.005', 'SGLD_thermal_0.0001/action_noise_0.1']\n",
    "Humanoid_1 = ['RMSprop/action_noise_0.01/alpha_0.1/', 'ExtraAdam/action_noise_0.01/alpha_0.1/', 'SGLD_thermal_1e-05/action_noise_0.1/alpha_0.1/']\n",
    "Humanoid_2 = ['RMSprop/action_noise_0.01/alpha_0.2/', 'ExtraAdam/action_noise_0.01/alpha_0.2/', 'SGLD_thermal_0.001/action_noise_0.005/alpha_0.2/']\n",
    "\n",
    "Inv_0 = ['RMSprop/action_noise_0.01', 'ExtraAdam/action_noise_0.005', 'SGLD_thermal_0.0001/action_noise_0.01']\n",
    "Inv_1 = ['RMSprop/action_noise_0.01/alpha_0.1/', 'ExtraAdam/action_noise_0.01/alpha_0.1/','SGLD_thermal_0.001/action_noise_0.01/alpha_0.1/']\n",
    "Inv_2 = ['RMSprop/action_noise_0.005/alpha_0.2/', 'ExtraAdam/action_noise_0.01/alpha_0.2/','SGLD_thermal_0.0001/action_noise_0.01/alpha_0.2/']\n",
    "\n",
    "params = [Walker_0, Walker_1, Walker_2, HalfCheetah_0, HalfCheetah_1, HalfCheetah_2, Humanoid_0, Humanoid_1, Humanoid_2, Inv_0, Inv_1, Inv_2]   \n",
    "plot_learning_all(params, title = 'Comparison', comparison = True, OnePlayer = False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Mass / Friction Uncertainty"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 101,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_env(env_name, params, ax, var, OnePlayer, eval_type):\n",
    "\n",
    "    if OnePlayer:\n",
    "        base_dir = os.getcwd() + '/models_OnePlayer/' + env_name\n",
    "    else:\n",
    "        base_dir = os.getcwd() + '/models_TwoPlayer/' + env_name\n",
    "    results = {}\n",
    "        \n",
    "    for i, exp in enumerate(params):\n",
    "\n",
    "        results[exp] = {}\n",
    "        for s in range(1):\n",
    "            try:        \n",
    "                if(OnePlayer):\n",
    "                    with open(base_dir + '/' + exp + '/alpha_0.0/' + str(s) + '/' + eval_type , 'rb') as f:\n",
    "                        results[exp][s] = pickle.load(f)\n",
    "                else:\n",
    "                    with open(base_dir + '/' + exp + str(s) + '/' + eval_type, '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, eval_type)\n",
    "\n",
    "def plot_model(results, ax, env_name, var, OnePlayer, eval_type):\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_friction']:\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(\"alpha_0.2\" in exp):\n",
    "            alpha = 0.2\n",
    "        elif(\"alpha_0.1\" in exp):\n",
    "            alpha = 0.1\n",
    "        else:\n",
    "            alpha = 0\n",
    "\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",
    "\n",
    "    ax.set_title(env_name + r' ($\\delta=$'+ str(alpha) + \")\")         \n",
    "    ax.set_xticks([0, 5, 10, 15, 20])\n",
    "    ax.set_xticklabels([0.5, 0.75, 1.0, 1.5, 2.0])\n",
    "    if(eval_type == 'results_model'):\n",
    "        ax.set_xlabel('Relative Mass')\n",
    "    else:\n",
    "        ax.set_xlabel('Relative Friction')\n",
    "    ax.set_ylabel('Reward')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 104,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_env_all(params, title, comparison, OnePlayer):\n",
    "    cols = 3\n",
    "    rows = 4\n",
    "    eval_type = ['results_model', 'results_friction']\n",
    "    envs = ['Walker2d-v2', 'Walker2d-v2', 'Walker2d-v2', \\\n",
    "            'HalfCheetah-v2', 'HalfCheetah-v2', 'HalfCheetah-v2', \\\n",
    "            'Humanoid-v2', 'Humanoid-v2', 'Humanoid-v2',\\\n",
    "            'InvertedPendulum-v2', 'InvertedPendulum-v2', 'InvertedPendulum-v2'\n",
    "            ]\n",
    "    for uncertainty in eval_type:\n",
    "        fig, axs = plt.subplots(rows, cols, figsize=(20*cols/3, 10*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",
    "\n",
    "            if(i%3 == 0):\n",
    "                plot_env(env, param, ax, var = True, OnePlayer = True, eval_type = uncertainty)   \n",
    "            else:\n",
    "                plot_env(env, param, ax, var = True, OnePlayer = False, eval_type = uncertainty)   \n",
    "            if(comparison):\n",
    "                ax.legend(loc='lower left')\n",
    "            else:\n",
    "                if i == 0:\n",
    "                    ax.legend()\n",
    "        plt.tight_layout()\n",
    "        \n",
    "        if(uncertainty == 'results_model'):\n",
    "            path = 'plots/Mass_Uncertainty_' + title \n",
    "        else:\n",
    "            path = 'plots/Friction_Uncertainty_' + title    \n",
    "            \n",
    "        plt.savefig(path + '.pdf')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 105,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 452268 characters]\n",
      "text/plain": [
       "<Figure size 1440x1440 with 12 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 400256 characters]\n",
      "text/plain": [
       "<Figure size 1440x1440 with 12 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "Walker_0 = ['RMSprop/action_noise_0.01', 'ExtraAdam/action_noise_0.1', 'SGLD_thermal_1e-05/action_noise_0.01']\n",
    "Walker_1 = ['RMSprop/action_noise_0.01/alpha_0.1/', 'ExtraAdam/action_noise_0.01/alpha_0.1/','SGLD_thermal_0.001/action_noise_0.01/alpha_0.1/']\n",
    "Walker_2 = ['RMSprop/action_noise_0.005/alpha_0.2/', 'ExtraAdam/action_noise_0.005/alpha_0.2/','SGLD_thermal_0.0001/action_noise_0.005/alpha_0.2/']\n",
    "\n",
    "HalfCheetah_0 = ['RMSprop/action_noise_0.01', 'ExtraAdam/action_noise_0.001', 'SGLD_thermal_1e-05/action_noise_0.01']\n",
    "HalfCheetah_1 = ['RMSprop/action_noise_0.01/alpha_0.1/', 'ExtraAdam/action_noise_0.01/alpha_0.1/', 'SGLD_thermal_1e-05/action_noise_0.1/alpha_0.1/']\n",
    "HalfCheetah_2 = ['RMSprop/action_noise_0.005/alpha_0.2/', 'ExtraAdam/action_noise_0.1/alpha_0.2/', 'SGLD_thermal_1e-05/action_noise_0.005/alpha_0.2/']  \n",
    "\n",
    "Humanoid_0 = ['RMSprop/action_noise_0.1', 'ExtraAdam/action_noise_0.005', 'SGLD_thermal_0.0001/action_noise_0.1']\n",
    "Humanoid_1 = ['RMSprop/action_noise_0.01/alpha_0.1/', 'ExtraAdam/action_noise_0.01/alpha_0.1/', 'SGLD_thermal_1e-05/action_noise_0.1/alpha_0.1/']\n",
    "Humanoid_2 = ['RMSprop/action_noise_0.01/alpha_0.2/', 'ExtraAdam/action_noise_0.01/alpha_0.2/', 'SGLD_thermal_0.001/action_noise_0.005/alpha_0.2/']\n",
    "\n",
    "Inv_0 = ['RMSprop/action_noise_0.01', 'ExtraAdam/action_noise_0.005', 'SGLD_thermal_0.0001/action_noise_0.01']\n",
    "Inv_1 = ['RMSprop/action_noise_0.01/alpha_0.1/', 'ExtraAdam/action_noise_0.01/alpha_0.1/','SGLD_thermal_0.001/action_noise_0.01/alpha_0.1/']\n",
    "Inv_2 = ['RMSprop/action_noise_0.005/alpha_0.2/', 'ExtraAdam/action_noise_0.01/alpha_0.2/','SGLD_thermal_0.0001/action_noise_0.01/alpha_0.2/']\n",
    "\n",
    "params = [Walker_0, Walker_1, Walker_2, HalfCheetah_0, HalfCheetah_1, HalfCheetah_2, Humanoid_0, Humanoid_1, Humanoid_2, Inv_0, Inv_1, Inv_2]   \n",
    "\n",
    "plot_env_all(params, title = 'Comparison', comparison = True, OnePlayer = False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Heat Map"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 107,
   "metadata": {},
   "outputs": [],
   "source": [
    "from mpl_toolkits.axes_grid1 import make_axes_locatable\n",
    "\n",
    "def plot_2d(env_name, params, OnePlayer):\n",
    "    if OnePlayer:\n",
    "        base_dir = os.getcwd() + '/models_OnePlayer/' + env_name\n",
    "    else:\n",
    "        base_dir = os.getcwd() + '/models_TwoPlayer/' + env_name\n",
    "\n",
    "    results = {}\n",
    "    vmin = 10000\n",
    "    vmax = -10000\n",
    "    for idx, param in enumerate(params):\n",
    "        try:\n",
    "            mat = np.zeros((11, 11))\n",
    "            for s in range(1):\n",
    "                if(OnePlayer):\n",
    "                    with open(base_dir + '/' + param + '/alpha_0.0/' +str(s) +'/results_friction_noise', 'rb') as f:\n",
    "                        data = pickle.load(f)\n",
    "                else:\n",
    "                    with open(base_dir + '/' + param + '/alpha_0.1/'+str(s) +'/results_friction_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\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=(10, 10))\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",
    "#         labels = [0.5, 0.7, 0.9, 1.2, 1.6, 2.0]\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",
    "        plt.ylabel('Relative Friction', 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",
    "#     plt.savefig(env_name +'_model_friction.pdf', bbox_inches='tight',pad_inches=0.0, dpi= 300)\n",
    "\n",
    "    \n",
    "#     if(OnePlayer):\n",
    "#         folder = 'OnePlayer'\n",
    "#     else:\n",
    "#         folder = 'TwoPlayer'\n",
    "    \n",
    "\n",
    "#     plt.savefig('plots/'+ env_name + '/' + folder + '/Heatmap/friction_noise.pdf', bbox_inches='tight',pad_inches=0.0, dpi = 300) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 108,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 9584 characters]\n",
      "text/plain": [
       "<Figure size 720x720 with 6 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "HalfCheetah_exp = ['RMSprop/action_noise_0.01', 'ExtraAdam/action_noise_0.001', 'SGLD_thermal_0.001/action_noise_0.005']\n",
    "plot_2d('HalfCheetah-v2', HalfCheetah_exp, OnePlayer = False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 109,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 9552 characters]\n",
      "text/plain": [
       "<Figure size 720x720 with 6 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "## Comparison One Player Average\n",
    "HalfCheetah_exp = ['RMSprop/action_noise_0.01', 'ExtraAdam/action_noise_0.001', 'SGLD_thermal_0.001/action_noise_0.01']\n",
    "plot_2d('HalfCheetah-v2', HalfCheetah_exp, OnePlayer = True)"
   ]
  },
  {
   "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.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
