{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib\n",
    "import matplotlib.pyplot as plt\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "import argparse\n",
    "import os\n",
    "import pickle\n",
    "from os import path\n",
    "import shutil\n",
    "\n",
    "import matplotlib as mpl\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import seaborn as sns\n",
    "import wandb"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Hyper-parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "hyperparameters = []\n",
    "smooth_weight = 0.90\n",
    "\n",
    "# run name format is <env_case>_<exp_type>_<algo_name>_<launch_time>\n",
    "reset_dir = False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "wandb_project = \"<wandb/project/name>\"\n",
    "dir_name = \"<project/directory>\"\n",
    "\n",
    "smooth_features_of_interest = [\"fair/fair_gap\", \"fair/fair_return\", \"fair/0/avg_return\", \"fair/1/avg_return\", \"fair/2/avg_return\", ] \n",
    "running_features_of_interest = [\"queue/fair_gap\", \"queue/fair_return\", \"queue/0/avg_return\", \"queue/1/avg_return\", \"queue/2/avg_return\", ]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Plotting specific parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 80,
   "metadata": {},
   "outputs": [],
   "source": [
    "axis_label_font_size = 24 #18\n",
    "axis_title_font_size = 28 #20\n",
    "plt_title_font_size = 24\n",
    "\n",
    "x_label = \"Time-steps\" # the label of x-axis\n",
    "y_label = \"Episodic Return\" #the label of y-axis\n",
    "y_lim_bottom = 0 #the bottom limit for the y-axis for fairness gap\n",
    "output_format = 'pdf' #`pdf`, `png`, or `svg`\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "metadata": {},
   "outputs": [],
   "source": [
    "env_dict = {\n",
    "    \"cheetah-family\": \"Half-Cheetah Variations : Big-Feet vs High-Friction vs Default Half-Cheetah-v3\",\n",
    "}\n",
    "\n",
    "algo_dict = {\n",
    "    \"lag-ppo\": \"Lagrangian-PPO\",\n",
    "    \"fcpo\": \"FOC-PPO\",\n",
    "    \"ppo\": \"PPO\",\n",
    "    \"m-ppo\": \"Minimal-PPO\",\n",
    "}\n",
    "\n",
    "metric_dict = {\n",
    "    \"fair/0/avg_return\":\"Epsiodic return\", \n",
    "    \"fair/1/avg_return\":\"Episodic return\",\n",
    "    \"fair/2/avg_return\":\"Episodic return\",\n",
    "    \"fair/fair_gap\": \"Difference in returns\", \n",
    "    \"fair/fair_return\": \"Cumulative return\", \n",
    "    \"queue/0/avg_return\":\"Epsiodic return\", \n",
    "    \"queue/1/avg_return\":\"Episodic return\",\n",
    "    \"queue/2/avg_return\":\"Episodic return\",\n",
    "    \"queue/fair_gap\": \"Difference in returns\", \n",
    "    \"queue/fair_return\": \"Cumulative return\"\n",
    "}\n",
    "\n",
    "title_dict = {\n",
    "    \"fair/0/avg_return\":\"Big-feet subgroup\",\n",
    "    \"fair/1/avg_return\":\"High-friction subgroup\",\n",
    "    \"fair/2/avg_return\":\"Default subgroup\",\n",
    "    \"fair/fair_gap\": \"Fairness gap, $\\epsilon$\", \n",
    "    \"fair/fair_return\": \"All subgroups\", \n",
    "    \"queue/0/avg_return\":\"Big-feet subgroup\",\n",
    "    \"queue/1/avg_return\":\"High-friction subgroup\",\n",
    "    \"queue/2/avg_return\":\"Default subgroup\",\n",
    "    \"queue/fair_gap\": \"Fairness gap, $\\epsilon$\", \n",
    "    \"queue/fair_return\": \"All subgroups\"\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Get data from WandB"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "api = wandb.Api(timeout=200)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "def all_features_present(summary):\n",
    "    \"\"\"\n",
    "    checks if all metrics are present in the run\n",
    "    :param summary:\n",
    "    :return:\n",
    "    \"\"\"\n",
    "    okay = True\n",
    "    for feature in smooth_features_of_interest + running_features_of_interest + [\"global_step\"]:\n",
    "        okay = okay and (feature in summary)\n",
    "    return okay"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "# delete the cached runs if specified\n",
    "if reset_dir:\n",
    "    try:\n",
    "        shutil.rmtree(dir_name)\n",
    "    except OSError as e:\n",
    "        print(\"Error: %s - %s.\" % (e.filename, e.strerror))\n",
    "\n",
    "# else load from the cache\n",
    "if not os.path.exists(dir_name):\n",
    "    os.makedirs(dir_name)\n",
    "    with open(f\"{dir_name}/cache.pkl\", \"wb\") as handle:\n",
    "        pickle.dump([[], [], [], {}, [], set([])], handle, protocol=pickle.HIGHEST_PROTOCOL)\n",
    "with open(f\"{dir_name}/cache.pkl\", \"rb\") as handle:\n",
    "    summary_list, config_list, name_list, envs, exp_names, ids = pickle.load(handle)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "data loaded\n"
     ]
    }
   ],
   "source": [
    "# Get all the runs in the project\n",
    "runs = api.runs(wandb_project)\n",
    "data = []\n",
    "\n",
    "for idx, run in enumerate(runs):\n",
    "    \n",
    "    # if the current run is not saved already\n",
    "    if run.id not in ids and run.config[\"env_case\"]==\"cheetah-family\":\n",
    "        \n",
    "        ids.add(run.id) # update the saved run history\n",
    "        print(f\"Adding run {run.id}\")\n",
    "\n",
    "        if all_features_present(run.summary) and run.state == \"finished\":\n",
    "            # Note: WandB api is shaky, the samples we get for each are different in run.scan_history\n",
    "            metric_history = run.scan_history(keys=smooth_features_of_interest + running_features_of_interest + [\"global_step\"])\n",
    "            metric_data = [(row[smooth_features_of_interest[0]], row[smooth_features_of_interest[1]],\n",
    "                            row[smooth_features_of_interest[2]], row[smooth_features_of_interest[3]], row[smooth_features_of_interest[4]],\n",
    "                            row[running_features_of_interest[0]], row[running_features_of_interest[1]],\n",
    "                            row[running_features_of_interest[2]], row[running_features_of_interest[3]], row[running_features_of_interest[4]],\n",
    "                            row[\"global_step\"]) for row in metric_history]\n",
    "            \n",
    "            metrics_dataframe = pd.DataFrame(metric_data, columns=smooth_features_of_interest + running_features_of_interest + [\"global_step\"])\n",
    "            \n",
    "            if metrics_dataframe.empty:\n",
    "                raise Exception(f\"Something unexpected hapened with run:{run.id}, with {run.config['exp_name']} and seed {run.config['seed']}\")\n",
    "            \n",
    "            # data collected!\n",
    "            # exp_name = <exp_type>_<algo>\n",
    "            # where exp_type = <exp_subname>_<eps> for fair algorithms and just <exp_subname> for unfair algos\n",
    "            exp_name = run.config[\"exp_name\"]\n",
    "\n",
    "            # get the algo and add it as a column\n",
    "            algo_name = exp_name.split(\"_\")[-1]\n",
    "            metrics_dataframe.insert(len(metrics_dataframe.columns), \"algo\", algo_name)\n",
    "\n",
    "            # get runs for all corresponding params tuned (eg nu_lr)\n",
    "            for param in hyperparameters:\n",
    "                if param in run.config:\n",
    "                    # update the exp_name to reflect the hyperparam value\n",
    "                    exp_name += \"-\" + param + \"-\" + str(run.config[param]) + \"-\"\n",
    "                    # creates a separate entry for each hyperparameter\n",
    "            metrics_dataframe.insert(len(metrics_dataframe.columns), \"exp_name\", exp_name)\n",
    "\n",
    "            # add the exp to the exp_name list\n",
    "            exp_names += [exp_name]\n",
    "\n",
    "            # add the seed also as a column in df\n",
    "            metrics_dataframe.insert(len(metrics_dataframe.columns), \"seed\", run.config[\"seed\"])\n",
    "            \n",
    "            # add the df to the all data\n",
    "            data += [metrics_dataframe]\n",
    "\n",
    "            # if the env is not previously seen\n",
    "            if run.config[\"env_case\"] not in envs:\n",
    "                # add to the envs dict\n",
    "                envs[run.config[\"env_case\"]] = [metrics_dataframe]\n",
    "            else:\n",
    "                # else, append the current data to the already present env log\n",
    "                envs[run.config[\"env_case\"]] += [metrics_dataframe]\n",
    "\n",
    "            # update the summary_list, config_list and name_list\n",
    "            # run.summary are the output key/values like accuracy.  We call ._json_dict to omit large files\n",
    "            summary_list.append(run.summary._json_dict)\n",
    "\n",
    "            # run.config is the input metrics.  We remove special values that start with _.\n",
    "            config_list.append({k: v for k, v in run.config.items() if not k.startswith(\"_\")})\n",
    "\n",
    "            # run.name is the name of the run.\n",
    "            name_list.append(run.name)\n",
    "\n",
    "# Convert all the df to a unified df\n",
    "summary_df = pd.DataFrame.from_records(summary_list)\n",
    "config_df = pd.DataFrame.from_records(config_list)\n",
    "name_df = pd.DataFrame({\"name\": name_list})\n",
    "all_df = pd.concat([name_df, config_df, summary_df], axis=1)\n",
    "# data = pd.concat(data, ignore_index=True)\n",
    "with open(f\"{dir_name}/cache.pkl\", \"wb\") as handle:\n",
    "    pickle.dump([summary_list, config_list, name_list, envs, exp_names, ids], handle, protocol=pickle.HIGHEST_PROTOCOL)\n",
    "print(\"data loaded\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Plotting"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "# create paths for storing plots\n",
    "if not os.path.exists(f\"{dir_name}/data/smooth\"):\n",
    "    os.makedirs(f\"{dir_name}/data/smooth\")\n",
    "if not os.path.exists(f\"{dir_name}/data/running\"):\n",
    "    os.makedirs(f\"{dir_name}/data/running\")\n",
    "\n",
    "if not os.path.exists(f\"{dir_name}/plots/smooth\"):\n",
    "    os.makedirs(f\"{dir_name}/plots/smooth\")\n",
    "if not os.path.exists(f\"{dir_name}/plots/running\"):\n",
    "    os.makedirs(f\"{dir_name}/plots/running\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Color schemes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 76,
   "metadata": {},
   "outputs": [],
   "source": [
    "# sns.set(style=\"darkgrid\")\n",
    "sns.set(style=\"white\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'lag-ppo': (0.2980392156862745, 0.4470588235294118, 0.6901960784313725),\n",
       " 'fcpo': (0.8666666666666667, 0.5176470588235295, 0.3215686274509804),\n",
       " 'ppo': (0.3333333333333333, 0.6588235294117647, 0.40784313725490196),\n",
       " 'm-ppo': (0.7686274509803922, 0.3058823529411765, 0.3215686274509804)}"
      ]
     },
     "execution_count": 43,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "current_palette = sns.color_palette(n_colors=len(algo_dict))\n",
    "current_palette_dict = dict(zip(algo_dict.keys(), current_palette))\n",
    "current_palette_dict"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Plotting utilities"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_interested_exp_names(name, epsilon, skip_minimal=False):\n",
    "    \"\"\"\n",
    "    exp_name format is:\n",
    "    for unfair algo: <name>_<algo>\n",
    "    for fair algo: <name>_eps_<epsilon>_<algo>\n",
    "    \"\"\"\n",
    "    interested_exp_names = []\n",
    "    for algo in algo_dict.keys():\n",
    "        if skip_minimal and algo==\"m-ppo\":\n",
    "            continue\n",
    "        \n",
    "        if algo == \"ppo\" or algo == \"m-ppo\":\n",
    "            exp_name = f\"{name}_{algo}\"\n",
    "        else:\n",
    "            exp_name = f\"{name}_eps_{epsilon}_{algo}\"\n",
    "        \n",
    "        assert exp_name in set(exp_names), f\"Exp {exp_name} not found!!\"\n",
    "\n",
    "        interested_exp_names += [exp_name]\n",
    "        \n",
    "    return sorted(interested_exp_names)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "# apply smoothing to the interested quantities\n",
    "\n",
    "# https://stackoverflow.com/questions/42281844/what-is-the-mathematics-behind-the-smoothing-parameter-in-tensorboards-scalar#_=_\n",
    "def smooth(scalars, weight):  # Weight between 0 and 1\n",
    "    last = scalars[0]  # First value in the plot (first timestep)\n",
    "    smoothed = list()\n",
    "    for point in scalars:\n",
    "        smoothed_val = last * weight + (1 - weight) * point  # Calculate smoothed value\n",
    "        smoothed.append(smoothed_val)  # Save it\n",
    "        last = smoothed_val  # Anchor the last smoothed value\n",
    "\n",
    "    return smoothed\n",
    "\n",
    "def _smooth(df_orig):\n",
    "    df = df_orig.copy()\n",
    "    for feature in smooth_features_of_interest:\n",
    "        df[feature] = smooth(list(df[feature]), smooth_weight)\n",
    "    return df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 83,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_metrics(df, eps, env, mode, save=True):\n",
    "    \"\"\"\n",
    "    the main plotting function\n",
    "    \"\"\"\n",
    "    width = 6 * 5\n",
    "    height = 6  * 1 \n",
    "    \n",
    "    if mode == \"smooth\":\n",
    "        features_of_interest = smooth_features_of_interest\n",
    "    elif mode == \"running\":\n",
    "        features_of_interest = running_features_of_interest\n",
    "    else:\n",
    "        raise Exception(\"unknown mode\")\n",
    "    \n",
    "    fig, ax = plt.subplots(1, 5, figsize=(width, height))\n",
    "    \n",
    "    for idx, feature in enumerate(features_of_interest): \n",
    "        sns.lineplot(data=df, x=\"global_step\", y=feature, hue=\"algo\", ci=\"sd\", palette=current_palette_dict, ax=ax[idx])\n",
    "        ax[idx].ticklabel_format(style=\"sci\", scilimits=(0, 0), axis=\"x\")\n",
    "        ax[idx].set(xlabel=x_label, ylabel=metric_dict[feature])\n",
    "        ax[idx].legend().remove()\n",
    "        ax[idx].set_title(title_dict[feature])\n",
    "        \n",
    "        ax[idx].title.set_size(axis_title_font_size)\n",
    "        ax[idx].xaxis.label.set_size(axis_label_font_size)\n",
    "        ax[idx].yaxis.label.set_size(axis_label_font_size)\n",
    "        \n",
    "    # for fairness gap indicate the threshold also\n",
    "    ax[0].axhline(y=eps,c=\"black\",linestyle=\"--\")\n",
    "    \n",
    "    # https://stackoverflow.com/a/45161551\n",
    "    plt.tight_layout(rect=[0, 0.03, 1, 0.95])\n",
    "\n",
    "    # https://stackoverflow.com/a/46267860\n",
    "    fig.subplots_adjust(bottom=0.3, wspace=0.33)\n",
    "    handles, labels = ax[0].get_legend_handles_labels()\n",
    "    # remove the first redundant lable\n",
    "    handles = handles[1:]\n",
    "    labels = [algo_dict[label] for label in labels[1:]]\n",
    "    \n",
    "    # https://stackoverflow.com/a/9764364\n",
    "    labels, handles = (list(t) for t in zip(*sorted(zip(labels, handles))))\n",
    "    \n",
    "    ax[1].legend(handles = handles , labels= labels, prop={'size':20}, \n",
    "                 loc='upper center', \n",
    "                 bbox_to_anchor=(0., -0.2),fancybox=False, shadow=False, ncol=4)\n",
    "    \n",
    "    if save:\n",
    "        plt.savefig(f\"{dir_name}/plots/{mode}/{env}_eps_{eps}.{output_format}\")\n",
    "    plt.show()\n",
    "    plt.clf()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'cheetah-family'}"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "set(all_df[\"env_case\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 65,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([  500, 10000,  1000])"
      ]
     },
     "execution_count": 65,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "config_df.epsilon.unique()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Loop over envs and plot"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [],
   "source": [
    "redo_smooth = True\n",
    "\n",
    "eps_list = config_df.epsilon.unique()\n",
    "int_exp_suffix = \"cheetah_final\" # plot corresponding to this experiment name"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Plotting smoothed plots"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 85,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "cheetah-family, eps=500, seeds=10\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 371528 characters]\n",
      "text/plain": [
       "<Figure size 2160x432 with 5 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "cheetah-family, eps=10000, seeds=10\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<Figure size 432x288 with 0 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 314644 characters]\n",
      "text/plain": [
       "<Figure size 2160x432 with 5 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "cheetah-family, eps=1000, seeds=10\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<Figure size 432x288 with 0 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 372004 characters]\n",
      "text/plain": [
       "<Figure size 2160x432 with 5 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "<Figure size 432x288 with 0 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "for env in set(all_df[\"env_case\"]):\n",
    "    \n",
    "    data = None\n",
    "    \n",
    "    # Load data for env\n",
    "    if redo_smooth or not path.exists(f\"{dir_name}/data/smooth/{env}.pkl\"):\n",
    "        with open(f\"{dir_name}/data/smooth/{env}.pkl\", \"wb\") as handle:\n",
    "            # smooth the indvidual runs\n",
    "            smoothed_dfs = []\n",
    "            for df in envs[env]:\n",
    "                smoothed_dfs += [_smooth(df)] \n",
    "            data = pd.concat(smoothed_dfs, sort=True)   \n",
    "            # Alternate option, can get data without smoothing and perform smoothing later \n",
    "            # data = pd.concat(envs[env], sort=True)  \n",
    "            assert data.shape[0] == envs[env][0].shape[0] * len(envs[env])\n",
    "            data[\"seed\"] = data[\"seed\"].astype(float)\n",
    "            for feature in smooth_features_of_interest:\n",
    "                data[feature] = data[feature].astype(float)\n",
    "            pickle.dump(data, handle, protocol=pickle.HIGHEST_PROTOCOL)\n",
    "    else:\n",
    "        with open(f\"{dir_name}/data/smooth/{env}.pkl\", \"rb\") as handle:\n",
    "            data = pickle.load(handle)\n",
    "        print(f\"{env}'s data loaded from cache\")\n",
    "    \n",
    "    # ----------------------------------------------------------------\n",
    "    #               Actual plotting calls here\n",
    "    # ----------------------------------------------------------------\n",
    "    \n",
    "    for eps in eps_list:\n",
    "        \n",
    "        # get interested exp names \n",
    "        interested_exp_names =  get_interested_exp_names(int_exp_suffix, eps, skip_minimal=False)\n",
    "    \n",
    "        # plot for one <env>_<eps> combination\n",
    "        plot_data = data.groupby([\"seed\", \"algo\"]).apply(lambda x: x)\n",
    "        plot_data = plot_data.loc[plot_data[\"exp_name\"].isin(interested_exp_names)]\n",
    "        \n",
    "        print(f\"{env}, eps={eps}, seeds={len(plot_data.seed.unique())}\")\n",
    "\n",
    "        # plot \n",
    "        plot_metrics(plot_data, eps, env, mode=\"smooth\", save=True)\n",
    "        \n",
    "        "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Plotting Running Plots"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 86,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "cheetah-family's data loaded from cache\n",
      "cheetah-family, eps=500, seeds=10\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 253492 characters]\n",
      "text/plain": [
       "<Figure size 2160x432 with 5 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "cheetah-family, eps=10000, seeds=10\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<Figure size 432x288 with 0 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 228764 characters]\n",
      "text/plain": [
       "<Figure size 2160x432 with 5 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "cheetah-family, eps=1000, seeds=10\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<Figure size 432x288 with 0 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 254400 characters]\n",
      "text/plain": [
       "<Figure size 2160x432 with 5 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "<Figure size 432x288 with 0 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "for env in set(all_df[\"env_case\"]):\n",
    "    \n",
    "    data = None\n",
    "    \n",
    "    # Load data for env\n",
    "    if not path.exists(f\"{dir_name}/data/running/{env}.pkl\"):\n",
    "        with open(f\"{dir_name}/data/running/{env}.pkl\", \"wb\") as handle:\n",
    "            data = pd.concat(envs[env], sort=True)  \n",
    "            assert data.shape[0] == envs[env][0].shape[0] * len(envs[env])\n",
    "            data[\"seed\"] = data[\"seed\"].astype(float)\n",
    "            for feature in running_features_of_interest:\n",
    "                data[feature] = data[feature].astype(float)\n",
    "            pickle.dump(data, handle, protocol=pickle.HIGHEST_PROTOCOL)\n",
    "    else:\n",
    "        with open(f\"{dir_name}/data/running/{env}.pkl\", \"rb\") as handle:\n",
    "            data = pickle.load(handle)\n",
    "        print(f\"{env}'s data loaded from cache\")\n",
    "    \n",
    "    # ----------------------------------------------------------------\n",
    "    #               Actual plotting calls here\n",
    "    # ----------------------------------------------------------------\n",
    "    \n",
    "    for eps in eps_list:\n",
    "        \n",
    "        # get interested exp names \n",
    "        interested_exp_names =  get_interested_exp_names(int_exp_suffix, eps, skip_minimal=False)\n",
    "    \n",
    "        # plot for one <env>_<eps> combination\n",
    "        plot_data = data.groupby([\"seed\", \"algo\"]).apply(lambda x: x)\n",
    "        plot_data = plot_data.loc[plot_data[\"exp_name\"].isin(interested_exp_names)]\n",
    "        \n",
    "        print(f\"{env}, eps={eps}, seeds={len(plot_data.seed.unique())}\")\n",
    "\n",
    "        # plot \n",
    "        plot_metrics(plot_data, eps, env, mode=\"running\", save=True)        \n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python [conda env:fair-rl] *",
   "language": "python",
   "name": "conda-env-fair-rl-py"
  },
  "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.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
