{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "8f4d6efb",
   "metadata": {},
   "source": [
    "# Results\n",
    "\n",
    "This notebook gathers results from evaluation JSON files and prints them as a list. \n",
    "\n",
    "### Setup\n",
    "\n",
    "- Set the root folder environment variable with `export CLIPORT_ROOT=<cliport_root>`\n",
    "- Train and evaluate agents by following the [README guide](https://github.com/cliport/cliport#single-task-training--evaluation)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "d072ae18",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "pybullet build time: Aug 16 2021 17:58:31\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "import sys\n",
    "import json\n",
    "\n",
    "from cliport import agents\n",
    "from cliport import tasks"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e2ee3b65",
   "metadata": {},
   "source": [
    "### Settings"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "95c14026",
   "metadata": {},
   "outputs": [],
   "source": [
    "root_folder = os.environ['CLIPORT_ROOT']\n",
    "exp_folder = os.path.join(root_folder, 'cliport_quickstart') # replace 'cliport_quickstart' with your exps folder"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2627285a",
   "metadata": {},
   "source": [
    "### Gather JSON Results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "6f5186e1",
   "metadata": {},
   "outputs": [],
   "source": [
    "tasks_list = list(tasks.names.keys())\n",
    "agents_list = list(agents.names.keys())\n",
    "demos_list = [1, 10, 100, 1000]\n",
    "\n",
    "results = {}\n",
    "for t in tasks_list:\n",
    "    for a in agents_list:\n",
    "        for d in demos_list:\n",
    "            task_folder = f'{t}-{a}-n{d}-train'\n",
    "            task_folder_path = os.path.join(exp_folder, task_folder, 'checkpoints')\n",
    "\n",
    "            if os.path.exists(task_folder_path):\n",
    "                jsons = [f for f in os.listdir(task_folder_path) if '.json' in f]\n",
    "                for j in jsons:\n",
    "                    model_type = 'multi' if 'multi' in j else 'single'\n",
    "                    eval_type = 'val' if 'val' in j else 'test'\n",
    "                    \n",
    "                    with open(os.path.join(task_folder_path, j)) as f:\n",
    "                        res = json.load(f)\n",
    "                    \n",
    "                    results[f'{t}-{a}-n{d}-{model_type}-{eval_type}'] = res"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0b6fcfa9",
   "metadata": {},
   "source": [
    "### Print Results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "2554998c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Experiments folder: /home/mshr/cliport/cliport_quickstart\n",
      "\n",
      "----- VAL -----\n",
      "\n",
      "stack-block-pyramid-seq-seen-colors | Train Demos: 1000\n",
      "\t97.3 : cliport | multi\n",
      "\n",
      "----- TEST -----\n",
      "\n",
      "stack-block-pyramid-seq-seen-colors | Train Demos: 1000\n",
      "\t96.5 : cliport | multi\n",
      "\n"
     ]
    }
   ],
   "source": [
    "print(f'Experiments folder: {exp_folder}\\n')\n",
    "\n",
    "for eval_type in ['val', 'test']:\n",
    "    print(f'----- {eval_type.upper()} -----\\n')\n",
    "    for t in tasks_list:\n",
    "        for a in agents_list:\n",
    "            for d in demos_list:\n",
    "                for model_type in ['single', 'multi']:\n",
    "                    eval_key = f'{t}-{a}-n{d}-{model_type}-{eval_type}'\n",
    "                    \n",
    "                    if eval_key in results:    \n",
    "                        print(f'{t} | Train Demos: {d}')\n",
    "                        \n",
    "                        res = results[eval_key]\n",
    "                        best_score, best_ckpt = max(zip([v['mean_reward'] for v in list(res.values())], \n",
    "                                                        res.keys())) # TODO: test that this works for full results folder\n",
    "                        \n",
    "                        print(f'\\t{best_score*100:1.1f} : {a} | {model_type}\\n')\n",
    "                            "
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.8.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}