notebooks/avalon_results.sync.ipynb

{
 "cells": [
  {
   "cell_type": "raw",
   "source": [
    "import json\n",
    "from collections import defaultdict\n",
    "from typing import Dict\n",
    "from typing import List\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import seaborn as sns\n",
    "from matplotlib import rcParams\n",
    "from matplotlib.colors import hsv_to_rgb\n",
    "from rliable import library as rly\n",
    "from rliable import metrics\n",
    "from tabulate import tabulate\n",
    "\n",
    "from avalon.common.log_utils import configure_local_logger\n",
    "from avalon.common.log_utils import logger\n",
    "from avalon.contrib.s3_utils import SimpleS3Client\n",
    "from avalon.datagen.world_creation.constants import TASKS_BY_TASK_GROUP\n",
    "from avalon.datagen.world_creation.constants import AvalonTaskGroup\n",
    "\n",
    "configure_local_logger()\n",
    "\n",
    "s3_client = SimpleS3Client(bucket_name=\"avalon-benchmark\")\n",
    "FIXED_WORLD_KEY = \"avalon_worlds__0824_full.tar.gz\"\n",
    "\n",
    "# to upload\n",
    "# data_key = \"avalon_worlds__d24d1746-4b0d-43da-bc83-338bc1eab441.tar.gz\"\n",
    "# s3_client.upload_from_file(Path(\"/tmp/avalon_worlds/d24d1746-4b0d-43da-bc83-338bc1eab441.tar.gz\"), data_key)"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": ""
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "ename": "EndpointConnectionError",
     "evalue": "Could not connect to the endpoint URL: \"https://avalon-benchmark.s3.us-west-2.amazonaws.com/random_policy__0__avalon_worlds__0824_full.tar.gz\"",
     "output_type": "error",
     "traceback": [
      "\u001B[0;31m---------------------------------------------------------------------------\u001B[0m",
      "\u001B[0;31mgaierror\u001B[0m                                  Traceback (most recent call last)",
      "File \u001B[0;32m/opt/venv/lib/python3.9/site-packages/urllib3/connection.py:174\u001B[0m, in \u001B[0;36mHTTPConnection._new_conn\u001B[0;34m(self)\u001B[0m\n\u001B[1;32m    173\u001B[0m \u001B[38;5;28;01mtry\u001B[39;00m:\n\u001B[0;32m--> 174\u001B[0m     conn \u001B[38;5;241m=\u001B[39m \u001B[43mconnection\u001B[49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43mcreate_connection\u001B[49m\u001B[43m(\u001B[49m\n\u001B[1;32m    175\u001B[0m \u001B[43m        \u001B[49m\u001B[43m(\u001B[49m\u001B[38;5;28;43mself\u001B[39;49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43m_dns_host\u001B[49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[38;5;28;43mself\u001B[39;49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43mport\u001B[49m\u001B[43m)\u001B[49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[38;5;28;43mself\u001B[39;49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43mtimeout\u001B[49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[38;5;241;43m*\u001B[39;49m\u001B[38;5;241;43m*\u001B[39;49m\u001B[43mextra_kw\u001B[49m\n\u001B[1;32m    176\u001B[0m \u001B[43m    \u001B[49m\u001B[43m)\u001B[49m\n\u001B[1;32m    178\u001B[0m \u001B[38;5;28;01mexcept\u001B[39;00m SocketTimeout:\n",
      "File \u001B[0;32m/opt/venv/lib/python3.9/site-packages/urllib3/util/connection.py:72\u001B[0m, in \u001B[0;36mcreate_connection\u001B[0;34m(address, timeout, source_address, socket_options)\u001B[0m\n\u001B[1;32m     68\u001B[0m     \u001B[38;5;28;01mreturn\u001B[39;00m six\u001B[38;5;241m.\u001B[39mraise_from(\n\u001B[1;32m     69\u001B[0m         LocationParseError(\u001B[38;5;124mu\u001B[39m\u001B[38;5;124m\"\u001B[39m\u001B[38;5;124m'\u001B[39m\u001B[38;5;132;01m%s\u001B[39;00m\u001B[38;5;124m'\u001B[39m\u001B[38;5;124m, label empty or too long\u001B[39m\u001B[38;5;124m\"\u001B[39m \u001B[38;5;241m%\u001B[39m host), \u001B[38;5;28;01mNone\u001B[39;00m\n\u001B[1;32m     70\u001B[0m     )\n\u001B[0;32m---> 72\u001B[0m \u001B[38;5;28;01mfor\u001B[39;00m res \u001B[38;5;129;01min\u001B[39;00m \u001B[43msocket\u001B[49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43mgetaddrinfo\u001B[49m\u001B[43m(\u001B[49m\u001B[43mhost\u001B[49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[43mport\u001B[49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[43mfamily\u001B[49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[43msocket\u001B[49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43mSOCK_STREAM\u001B[49m\u001B[43m)\u001B[49m:\n\u001B[1;32m     73\u001B[0m     af, socktype, proto, canonname, sa \u001B[38;5;241m=\u001B[39m res\n",
      "File \u001B[0;32m/usr/lib/python3.9/socket.py:953\u001B[0m, in \u001B[0;36mgetaddrinfo\u001B[0;34m(host, port, family, type, proto, flags)\u001B[0m\n\u001B[1;32m    952\u001B[0m addrlist \u001B[38;5;241m=\u001B[39m []\n\u001B[0;32m--> 953\u001B[0m \u001B[38;5;28;01mfor\u001B[39;00m res \u001B[38;5;129;01min\u001B[39;00m \u001B[43m_socket\u001B[49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43mgetaddrinfo\u001B[49m\u001B[43m(\u001B[49m\u001B[43mhost\u001B[49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[43mport\u001B[49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[43mfamily\u001B[49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[38;5;28;43mtype\u001B[39;49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[43mproto\u001B[49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[43mflags\u001B[49m\u001B[43m)\u001B[49m:\n\u001B[1;32m    954\u001B[0m     af, socktype, proto, canonname, sa \u001B[38;5;241m=\u001B[39m res\n",
      "\u001B[0;31mgaierror\u001B[0m: [Errno -3] Temporary failure in name resolution",
      "\nDuring handling of the above exception, another exception occurred:\n",
      "\u001B[0;31mNewConnectionError\u001B[0m                        Traceback (most recent call last)",
      "File \u001B[0;32m/opt/venv/lib/python3.9/site-packages/botocore/httpsession.py:455\u001B[0m, in \u001B[0;36mURLLib3Session.send\u001B[0;34m(self, request)\u001B[0m\n\u001B[1;32m    454\u001B[0m request_target \u001B[38;5;241m=\u001B[39m \u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39m_get_request_target(request\u001B[38;5;241m.\u001B[39murl, proxy_url)\n\u001B[0;32m--> 455\u001B[0m urllib_response \u001B[38;5;241m=\u001B[39m \u001B[43mconn\u001B[49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43murlopen\u001B[49m\u001B[43m(\u001B[49m\n\u001B[1;32m    456\u001B[0m \u001B[43m    \u001B[49m\u001B[43mmethod\u001B[49m\u001B[38;5;241;43m=\u001B[39;49m\u001B[43mrequest\u001B[49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43mmethod\u001B[49m\u001B[43m,\u001B[49m\n\u001B[1;32m    457\u001B[0m \u001B[43m    \u001B[49m\u001B[43murl\u001B[49m\u001B[38;5;241;43m=\u001B[39;49m\u001B[43mrequest_target\u001B[49m\u001B[43m,\u001B[49m\n\u001B[1;32m    458\u001B[0m \u001B[43m    \u001B[49m\u001B[43mbody\u001B[49m\u001B[38;5;241;43m=\u001B[39;49m\u001B[43mrequest\u001B[49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43mbody\u001B[49m\u001B[43m,\u001B[49m\n\u001B[1;32m    459\u001B[0m \u001B[43m    \u001B[49m\u001B[43mheaders\u001B[49m\u001B[38;5;241;43m=\u001B[39;49m\u001B[43mrequest\u001B[49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43mheaders\u001B[49m\u001B[43m,\u001B[49m\n\u001B[1;32m    460\u001B[0m \u001B[43m    \u001B[49m\u001B[43mretries\u001B[49m\u001B[38;5;241;43m=\u001B[39;49m\u001B[43mRetry\u001B[49m\u001B[43m(\u001B[49m\u001B[38;5;28;43;01mFalse\u001B[39;49;00m\u001B[43m)\u001B[49m\u001B[43m,\u001B[49m\n\u001B[1;32m    461\u001B[0m \u001B[43m    \u001B[49m\u001B[43massert_same_host\u001B[49m\u001B[38;5;241;43m=\u001B[39;49m\u001B[38;5;28;43;01mFalse\u001B[39;49;00m\u001B[43m,\u001B[49m\n\u001B[1;32m    462\u001B[0m \u001B[43m    \u001B[49m\u001B[43mpreload_content\u001B[49m\u001B[38;5;241;43m=\u001B[39;49m\u001B[38;5;28;43;01mFalse\u001B[39;49;00m\u001B[43m,\u001B[49m\n\u001B[1;32m    463\u001B[0m \u001B[43m    \u001B[49m\u001B[43mdecode_content\u001B[49m\u001B[38;5;241;43m=\u001B[39;49m\u001B[38;5;28;43;01mFalse\u001B[39;49;00m\u001B[43m,\u001B[49m\n\u001B[1;32m    464\u001B[0m \u001B[43m    \u001B[49m\u001B[43mchunked\u001B[49m\u001B[38;5;241;43m=\u001B[39;49m\u001B[38;5;28;43mself\u001B[39;49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43m_chunked\u001B[49m\u001B[43m(\u001B[49m\u001B[43mrequest\u001B[49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43mheaders\u001B[49m\u001B[43m)\u001B[49m\u001B[43m,\u001B[49m\n\u001B[1;32m    465\u001B[0m \u001B[43m\u001B[49m\u001B[43m)\u001B[49m\n\u001B[1;32m    467\u001B[0m http_response \u001B[38;5;241m=\u001B[39m botocore\u001B[38;5;241m.\u001B[39mawsrequest\u001B[38;5;241m.\u001B[39mAWSResponse(\n\u001B[1;32m    468\u001B[0m     request\u001B[38;5;241m.\u001B[39murl,\n\u001B[1;32m    469\u001B[0m     urllib_response\u001B[38;5;241m.\u001B[39mstatus,\n\u001B[1;32m    470\u001B[0m     urllib_response\u001B[38;5;241m.\u001B[39mheaders,\n\u001B[1;32m    471\u001B[0m     urllib_response,\n\u001B[1;32m    472\u001B[0m )\n",
      "File \u001B[0;32m/opt/venv/lib/python3.9/site-packages/urllib3/connectionpool.py:787\u001B[0m, in \u001B[0;36mHTTPConnectionPool.urlopen\u001B[0;34m(self, method, url, body, headers, retries, redirect, assert_same_host, timeout, pool_timeout, release_conn, chunked, body_pos, **response_kw)\u001B[0m\n\u001B[1;32m    785\u001B[0m     e \u001B[38;5;241m=\u001B[39m ProtocolError(\u001B[38;5;124m\"\u001B[39m\u001B[38;5;124mConnection aborted.\u001B[39m\u001B[38;5;124m\"\u001B[39m, e)\n\u001B[0;32m--> 787\u001B[0m retries \u001B[38;5;241m=\u001B[39m \u001B[43mretries\u001B[49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43mincrement\u001B[49m\u001B[43m(\u001B[49m\n\u001B[1;32m    788\u001B[0m \u001B[43m    \u001B[49m\u001B[43mmethod\u001B[49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[43murl\u001B[49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[43merror\u001B[49m\u001B[38;5;241;43m=\u001B[39;49m\u001B[43me\u001B[49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[43m_pool\u001B[49m\u001B[38;5;241;43m=\u001B[39;49m\u001B[38;5;28;43mself\u001B[39;49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[43m_stacktrace\u001B[49m\u001B[38;5;241;43m=\u001B[39;49m\u001B[43msys\u001B[49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43mexc_info\u001B[49m\u001B[43m(\u001B[49m\u001B[43m)\u001B[49m\u001B[43m[\u001B[49m\u001B[38;5;241;43m2\u001B[39;49m\u001B[43m]\u001B[49m\n\u001B[1;32m    789\u001B[0m \u001B[43m\u001B[49m\u001B[43m)\u001B[49m\n\u001B[1;32m    790\u001B[0m retries\u001B[38;5;241m.\u001B[39msleep()\n",
      "File \u001B[0;32m/opt/venv/lib/python3.9/site-packages/urllib3/util/retry.py:525\u001B[0m, in \u001B[0;36mRetry.increment\u001B[0;34m(self, method, url, response, error, _pool, _stacktrace)\u001B[0m\n\u001B[1;32m    523\u001B[0m \u001B[38;5;28;01mif\u001B[39;00m \u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39mtotal \u001B[38;5;129;01mis\u001B[39;00m \u001B[38;5;28;01mFalse\u001B[39;00m \u001B[38;5;129;01mand\u001B[39;00m error:\n\u001B[1;32m    524\u001B[0m     \u001B[38;5;66;03m# Disabled, indicate to re-raise the error.\u001B[39;00m\n\u001B[0;32m--> 525\u001B[0m     \u001B[38;5;28;01mraise\u001B[39;00m \u001B[43msix\u001B[49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43mreraise\u001B[49m\u001B[43m(\u001B[49m\u001B[38;5;28;43mtype\u001B[39;49m\u001B[43m(\u001B[49m\u001B[43merror\u001B[49m\u001B[43m)\u001B[49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[43merror\u001B[49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[43m_stacktrace\u001B[49m\u001B[43m)\u001B[49m\n\u001B[1;32m    527\u001B[0m total \u001B[38;5;241m=\u001B[39m \u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39mtotal\n",
      "File \u001B[0;32m/opt/venv/lib/python3.9/site-packages/urllib3/packages/six.py:770\u001B[0m, in \u001B[0;36mreraise\u001B[0;34m(tp, value, tb)\u001B[0m\n\u001B[1;32m    769\u001B[0m         \u001B[38;5;28;01mraise\u001B[39;00m value\u001B[38;5;241m.\u001B[39mwith_traceback(tb)\n\u001B[0;32m--> 770\u001B[0m     \u001B[38;5;28;01mraise\u001B[39;00m value\n\u001B[1;32m    771\u001B[0m \u001B[38;5;28;01mfinally\u001B[39;00m:\n",
      "File \u001B[0;32m/opt/venv/lib/python3.9/site-packages/urllib3/connectionpool.py:703\u001B[0m, in \u001B[0;36mHTTPConnectionPool.urlopen\u001B[0;34m(self, method, url, body, headers, retries, redirect, assert_same_host, timeout, pool_timeout, release_conn, chunked, body_pos, **response_kw)\u001B[0m\n\u001B[1;32m    702\u001B[0m \u001B[38;5;66;03m# Make the request on the httplib connection object.\u001B[39;00m\n\u001B[0;32m--> 703\u001B[0m httplib_response \u001B[38;5;241m=\u001B[39m \u001B[38;5;28;43mself\u001B[39;49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43m_make_request\u001B[49m\u001B[43m(\u001B[49m\n\u001B[1;32m    704\u001B[0m \u001B[43m    \u001B[49m\u001B[43mconn\u001B[49m\u001B[43m,\u001B[49m\n\u001B[1;32m    705\u001B[0m \u001B[43m    \u001B[49m\u001B[43mmethod\u001B[49m\u001B[43m,\u001B[49m\n\u001B[1;32m    706\u001B[0m \u001B[43m    \u001B[49m\u001B[43murl\u001B[49m\u001B[43m,\u001B[49m\n\u001B[1;32m    707\u001B[0m \u001B[43m    \u001B[49m\u001B[43mtimeout\u001B[49m\u001B[38;5;241;43m=\u001B[39;49m\u001B[43mtimeout_obj\u001B[49m\u001B[43m,\u001B[49m\n\u001B[1;32m    708\u001B[0m \u001B[43m    \u001B[49m\u001B[43mbody\u001B[49m\u001B[38;5;241;43m=\u001B[39;49m\u001B[43mbody\u001B[49m\u001B[43m,\u001B[49m\n\u001B[1;32m    709\u001B[0m \u001B[43m    \u001B[49m\u001B[43mheaders\u001B[49m\u001B[38;5;241;43m=\u001B[39;49m\u001B[43mheaders\u001B[49m\u001B[43m,\u001B[49m\n\u001B[1;32m    710\u001B[0m \u001B[43m    \u001B[49m\u001B[43mchunked\u001B[49m\u001B[38;5;241;43m=\u001B[39;49m\u001B[43mchunked\u001B[49m\u001B[43m,\u001B[49m\n\u001B[1;32m    711\u001B[0m \u001B[43m\u001B[49m\u001B[43m)\u001B[49m\n\u001B[1;32m    713\u001B[0m \u001B[38;5;66;03m# If we're going to release the connection in ``finally:``, then\u001B[39;00m\n\u001B[1;32m    714\u001B[0m \u001B[38;5;66;03m# the response doesn't need to know about the connection. Otherwise\u001B[39;00m\n\u001B[1;32m    715\u001B[0m \u001B[38;5;66;03m# it will also try to release it and we'll have a double-release\u001B[39;00m\n\u001B[1;32m    716\u001B[0m \u001B[38;5;66;03m# mess.\u001B[39;00m\n",
      "File \u001B[0;32m/opt/venv/lib/python3.9/site-packages/urllib3/connectionpool.py:386\u001B[0m, in \u001B[0;36mHTTPConnectionPool._make_request\u001B[0;34m(self, conn, method, url, timeout, chunked, **httplib_request_kw)\u001B[0m\n\u001B[1;32m    385\u001B[0m \u001B[38;5;28;01mtry\u001B[39;00m:\n\u001B[0;32m--> 386\u001B[0m     \u001B[38;5;28;43mself\u001B[39;49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43m_validate_conn\u001B[49m\u001B[43m(\u001B[49m\u001B[43mconn\u001B[49m\u001B[43m)\u001B[49m\n\u001B[1;32m    387\u001B[0m \u001B[38;5;28;01mexcept\u001B[39;00m (SocketTimeout, BaseSSLError) \u001B[38;5;28;01mas\u001B[39;00m e:\n\u001B[1;32m    388\u001B[0m     \u001B[38;5;66;03m# Py2 raises this as a BaseSSLError, Py3 raises it as socket timeout.\u001B[39;00m\n",
      "File \u001B[0;32m/opt/venv/lib/python3.9/site-packages/urllib3/connectionpool.py:1042\u001B[0m, in \u001B[0;36mHTTPSConnectionPool._validate_conn\u001B[0;34m(self, conn)\u001B[0m\n\u001B[1;32m   1041\u001B[0m \u001B[38;5;28;01mif\u001B[39;00m \u001B[38;5;129;01mnot\u001B[39;00m \u001B[38;5;28mgetattr\u001B[39m(conn, \u001B[38;5;124m\"\u001B[39m\u001B[38;5;124msock\u001B[39m\u001B[38;5;124m\"\u001B[39m, \u001B[38;5;28;01mNone\u001B[39;00m):  \u001B[38;5;66;03m# AppEngine might not have  `.sock`\u001B[39;00m\n\u001B[0;32m-> 1042\u001B[0m     \u001B[43mconn\u001B[49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43mconnect\u001B[49m\u001B[43m(\u001B[49m\u001B[43m)\u001B[49m\n\u001B[1;32m   1044\u001B[0m \u001B[38;5;28;01mif\u001B[39;00m \u001B[38;5;129;01mnot\u001B[39;00m conn\u001B[38;5;241m.\u001B[39mis_verified:\n",
      "File \u001B[0;32m/opt/venv/lib/python3.9/site-packages/urllib3/connection.py:358\u001B[0m, in \u001B[0;36mHTTPSConnection.connect\u001B[0;34m(self)\u001B[0m\n\u001B[1;32m    356\u001B[0m \u001B[38;5;28;01mdef\u001B[39;00m \u001B[38;5;21mconnect\u001B[39m(\u001B[38;5;28mself\u001B[39m):\n\u001B[1;32m    357\u001B[0m     \u001B[38;5;66;03m# Add certificate verification\u001B[39;00m\n\u001B[0;32m--> 358\u001B[0m     \u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39msock \u001B[38;5;241m=\u001B[39m conn \u001B[38;5;241m=\u001B[39m \u001B[38;5;28;43mself\u001B[39;49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43m_new_conn\u001B[49m\u001B[43m(\u001B[49m\u001B[43m)\u001B[49m\n\u001B[1;32m    359\u001B[0m     hostname \u001B[38;5;241m=\u001B[39m \u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39mhost\n",
      "File \u001B[0;32m/opt/venv/lib/python3.9/site-packages/urllib3/connection.py:186\u001B[0m, in \u001B[0;36mHTTPConnection._new_conn\u001B[0;34m(self)\u001B[0m\n\u001B[1;32m    185\u001B[0m \u001B[38;5;28;01mexcept\u001B[39;00m SocketError \u001B[38;5;28;01mas\u001B[39;00m e:\n\u001B[0;32m--> 186\u001B[0m     \u001B[38;5;28;01mraise\u001B[39;00m NewConnectionError(\n\u001B[1;32m    187\u001B[0m         \u001B[38;5;28mself\u001B[39m, \u001B[38;5;124m\"\u001B[39m\u001B[38;5;124mFailed to establish a new connection: \u001B[39m\u001B[38;5;132;01m%s\u001B[39;00m\u001B[38;5;124m\"\u001B[39m \u001B[38;5;241m%\u001B[39m e\n\u001B[1;32m    188\u001B[0m     )\n\u001B[1;32m    190\u001B[0m \u001B[38;5;28;01mreturn\u001B[39;00m conn\n",
      "\u001B[0;31mNewConnectionError\u001B[0m: <botocore.awsrequest.AWSHTTPSConnection object at 0x7f28c7401850>: Failed to establish a new connection: [Errno -3] Temporary failure in name resolution",
      "\nDuring handling of the above exception, another exception occurred:\n",
      "\u001B[0;31mEndpointConnectionError\u001B[0m                   Traceback (most recent call last)",
      "Cell \u001B[0;32mIn [2], line 6\u001B[0m\n\u001B[1;32m      4\u001B[0m combined_random_results \u001B[38;5;241m=\u001B[39m defaultdict(\u001B[38;5;28mlist\u001B[39m)\n\u001B[1;32m      5\u001B[0m \u001B[38;5;28;01mfor\u001B[39;00m result_key \u001B[38;5;129;01min\u001B[39;00m random_result_keys:\n\u001B[0;32m----> 6\u001B[0m     reported_data \u001B[38;5;241m=\u001B[39m json\u001B[38;5;241m.\u001B[39mloads(\u001B[43ms3_client\u001B[49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43mload\u001B[49m\u001B[43m(\u001B[49m\u001B[43mresult_key\u001B[49m\u001B[43m)\u001B[49m)\n\u001B[1;32m      7\u001B[0m     \u001B[38;5;28;01mfor\u001B[39;00m k, v \u001B[38;5;129;01min\u001B[39;00m reported_data[\u001B[38;5;124m\"\u001B[39m\u001B[38;5;124mall_results\u001B[39m\u001B[38;5;124m\"\u001B[39m]\u001B[38;5;241m.\u001B[39mitems():\n\u001B[1;32m      8\u001B[0m         combined_random_results[k]\u001B[38;5;241m.\u001B[39mappend(v)\n",
      "File \u001B[0;32m/opt/projects/avalon/avalon/contrib/s3_utils.py:52\u001B[0m, in \u001B[0;36mSimpleS3Client.load\u001B[0;34m(self, key)\u001B[0m\n\u001B[1;32m     50\u001B[0m \u001B[38;5;28;01mdef\u001B[39;00m \u001B[38;5;21mload\u001B[39m(\u001B[38;5;28mself\u001B[39m, key: \u001B[38;5;28mstr\u001B[39m) \u001B[38;5;241m-\u001B[39m\u001B[38;5;241m>\u001B[39m \u001B[38;5;28mbytes\u001B[39m:\n\u001B[1;32m     51\u001B[0m     \u001B[38;5;28;01mtry\u001B[39;00m:\n\u001B[0;32m---> 52\u001B[0m         data \u001B[38;5;241m=\u001B[39m \u001B[38;5;28;43mself\u001B[39;49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43mclient\u001B[49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43mget_object\u001B[49m\u001B[43m(\u001B[49m\u001B[43mBucket\u001B[49m\u001B[38;5;241;43m=\u001B[39;49m\u001B[38;5;28;43mself\u001B[39;49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43mbucket_name\u001B[49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[43mKey\u001B[49m\u001B[38;5;241;43m=\u001B[39;49m\u001B[43mkey\u001B[49m\u001B[43m)\u001B[49m[\u001B[38;5;124m\"\u001B[39m\u001B[38;5;124mBody\u001B[39m\u001B[38;5;124m\"\u001B[39m]\u001B[38;5;241m.\u001B[39mread()\n\u001B[1;32m     53\u001B[0m     \u001B[38;5;28;01mexcept\u001B[39;00m ClientError \u001B[38;5;28;01mas\u001B[39;00m e:\n\u001B[1;32m     54\u001B[0m         \u001B[38;5;28;01mif\u001B[39;00m e\u001B[38;5;241m.\u001B[39mresponse[\u001B[38;5;124m\"\u001B[39m\u001B[38;5;124mError\u001B[39m\u001B[38;5;124m\"\u001B[39m][\u001B[38;5;124m\"\u001B[39m\u001B[38;5;124mCode\u001B[39m\u001B[38;5;124m\"\u001B[39m] \u001B[38;5;241m==\u001B[39m \u001B[38;5;124m\"\u001B[39m\u001B[38;5;124mNoSuchKey\u001B[39m\u001B[38;5;124m\"\u001B[39m:\n",
      "File \u001B[0;32m/opt/venv/lib/python3.9/site-packages/botocore/client.py:514\u001B[0m, in \u001B[0;36mClientCreator._create_api_method.<locals>._api_call\u001B[0;34m(self, *args, **kwargs)\u001B[0m\n\u001B[1;32m    510\u001B[0m     \u001B[38;5;28;01mraise\u001B[39;00m \u001B[38;5;167;01mTypeError\u001B[39;00m(\n\u001B[1;32m    511\u001B[0m         \u001B[38;5;124mf\u001B[39m\u001B[38;5;124m\"\u001B[39m\u001B[38;5;132;01m{\u001B[39;00mpy_operation_name\u001B[38;5;132;01m}\u001B[39;00m\u001B[38;5;124m() only accepts keyword arguments.\u001B[39m\u001B[38;5;124m\"\u001B[39m\n\u001B[1;32m    512\u001B[0m     )\n\u001B[1;32m    513\u001B[0m \u001B[38;5;66;03m# The \"self\" in this scope is referring to the BaseClient.\u001B[39;00m\n\u001B[0;32m--> 514\u001B[0m \u001B[38;5;28;01mreturn\u001B[39;00m \u001B[38;5;28;43mself\u001B[39;49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43m_make_api_call\u001B[49m\u001B[43m(\u001B[49m\u001B[43moperation_name\u001B[49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[43mkwargs\u001B[49m\u001B[43m)\u001B[49m\n",
      "File \u001B[0;32m/opt/venv/lib/python3.9/site-packages/botocore/client.py:921\u001B[0m, in \u001B[0;36mBaseClient._make_api_call\u001B[0;34m(self, operation_name, api_params)\u001B[0m\n\u001B[1;32m    919\u001B[0m \u001B[38;5;28;01melse\u001B[39;00m:\n\u001B[1;32m    920\u001B[0m     apply_request_checksum(request_dict)\n\u001B[0;32m--> 921\u001B[0m     http, parsed_response \u001B[38;5;241m=\u001B[39m \u001B[38;5;28;43mself\u001B[39;49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43m_make_request\u001B[49m\u001B[43m(\u001B[49m\n\u001B[1;32m    922\u001B[0m \u001B[43m        \u001B[49m\u001B[43moperation_model\u001B[49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[43mrequest_dict\u001B[49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[43mrequest_context\u001B[49m\n\u001B[1;32m    923\u001B[0m \u001B[43m    \u001B[49m\u001B[43m)\u001B[49m\n\u001B[1;32m    925\u001B[0m \u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39mmeta\u001B[38;5;241m.\u001B[39mevents\u001B[38;5;241m.\u001B[39memit(\n\u001B[1;32m    926\u001B[0m     \u001B[38;5;124m'\u001B[39m\u001B[38;5;124mafter-call.\u001B[39m\u001B[38;5;132;01m{service_id}\u001B[39;00m\u001B[38;5;124m.\u001B[39m\u001B[38;5;132;01m{operation_name}\u001B[39;00m\u001B[38;5;124m'\u001B[39m\u001B[38;5;241m.\u001B[39mformat(\n\u001B[1;32m    927\u001B[0m         service_id\u001B[38;5;241m=\u001B[39mservice_id, operation_name\u001B[38;5;241m=\u001B[39moperation_name\n\u001B[0;32m   (...)\u001B[0m\n\u001B[1;32m    932\u001B[0m     context\u001B[38;5;241m=\u001B[39mrequest_context,\n\u001B[1;32m    933\u001B[0m )\n\u001B[1;32m    935\u001B[0m \u001B[38;5;28;01mif\u001B[39;00m http\u001B[38;5;241m.\u001B[39mstatus_code \u001B[38;5;241m>\u001B[39m\u001B[38;5;241m=\u001B[39m \u001B[38;5;241m300\u001B[39m:\n",
      "File \u001B[0;32m/opt/venv/lib/python3.9/site-packages/botocore/client.py:944\u001B[0m, in \u001B[0;36mBaseClient._make_request\u001B[0;34m(self, operation_model, request_dict, request_context)\u001B[0m\n\u001B[1;32m    942\u001B[0m \u001B[38;5;28;01mdef\u001B[39;00m \u001B[38;5;21m_make_request\u001B[39m(\u001B[38;5;28mself\u001B[39m, operation_model, request_dict, request_context):\n\u001B[1;32m    943\u001B[0m     \u001B[38;5;28;01mtry\u001B[39;00m:\n\u001B[0;32m--> 944\u001B[0m         \u001B[38;5;28;01mreturn\u001B[39;00m \u001B[38;5;28;43mself\u001B[39;49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43m_endpoint\u001B[49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43mmake_request\u001B[49m\u001B[43m(\u001B[49m\u001B[43moperation_model\u001B[49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[43mrequest_dict\u001B[49m\u001B[43m)\u001B[49m\n\u001B[1;32m    945\u001B[0m     \u001B[38;5;28;01mexcept\u001B[39;00m \u001B[38;5;167;01mException\u001B[39;00m \u001B[38;5;28;01mas\u001B[39;00m e:\n\u001B[1;32m    946\u001B[0m         \u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39mmeta\u001B[38;5;241m.\u001B[39mevents\u001B[38;5;241m.\u001B[39memit(\n\u001B[1;32m    947\u001B[0m             \u001B[38;5;124m'\u001B[39m\u001B[38;5;124mafter-call-error.\u001B[39m\u001B[38;5;132;01m{service_id}\u001B[39;00m\u001B[38;5;124m.\u001B[39m\u001B[38;5;132;01m{operation_name}\u001B[39;00m\u001B[38;5;124m'\u001B[39m\u001B[38;5;241m.\u001B[39mformat(\n\u001B[1;32m    948\u001B[0m                 service_id\u001B[38;5;241m=\u001B[39m\u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39m_service_model\u001B[38;5;241m.\u001B[39mservice_id\u001B[38;5;241m.\u001B[39mhyphenize(),\n\u001B[0;32m   (...)\u001B[0m\n\u001B[1;32m    952\u001B[0m             context\u001B[38;5;241m=\u001B[39mrequest_context,\n\u001B[1;32m    953\u001B[0m         )\n",
      "File \u001B[0;32m/opt/venv/lib/python3.9/site-packages/botocore/endpoint.py:119\u001B[0m, in \u001B[0;36mEndpoint.make_request\u001B[0;34m(self, operation_model, request_dict)\u001B[0m\n\u001B[1;32m    113\u001B[0m \u001B[38;5;28;01mdef\u001B[39;00m \u001B[38;5;21mmake_request\u001B[39m(\u001B[38;5;28mself\u001B[39m, operation_model, request_dict):\n\u001B[1;32m    114\u001B[0m     logger\u001B[38;5;241m.\u001B[39mdebug(\n\u001B[1;32m    115\u001B[0m         \u001B[38;5;124m\"\u001B[39m\u001B[38;5;124mMaking request for \u001B[39m\u001B[38;5;132;01m%s\u001B[39;00m\u001B[38;5;124m with params: \u001B[39m\u001B[38;5;132;01m%s\u001B[39;00m\u001B[38;5;124m\"\u001B[39m,\n\u001B[1;32m    116\u001B[0m         operation_model,\n\u001B[1;32m    117\u001B[0m         request_dict,\n\u001B[1;32m    118\u001B[0m     )\n\u001B[0;32m--> 119\u001B[0m     \u001B[38;5;28;01mreturn\u001B[39;00m \u001B[38;5;28;43mself\u001B[39;49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43m_send_request\u001B[49m\u001B[43m(\u001B[49m\u001B[43mrequest_dict\u001B[49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[43moperation_model\u001B[49m\u001B[43m)\u001B[49m\n",
      "File \u001B[0;32m/opt/venv/lib/python3.9/site-packages/botocore/endpoint.py:202\u001B[0m, in \u001B[0;36mEndpoint._send_request\u001B[0;34m(self, request_dict, operation_model)\u001B[0m\n\u001B[1;32m    198\u001B[0m request \u001B[38;5;241m=\u001B[39m \u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39mcreate_request(request_dict, operation_model)\n\u001B[1;32m    199\u001B[0m success_response, exception \u001B[38;5;241m=\u001B[39m \u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39m_get_response(\n\u001B[1;32m    200\u001B[0m     request, operation_model, context\n\u001B[1;32m    201\u001B[0m )\n\u001B[0;32m--> 202\u001B[0m \u001B[38;5;28;01mwhile\u001B[39;00m \u001B[38;5;28;43mself\u001B[39;49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43m_needs_retry\u001B[49m\u001B[43m(\u001B[49m\n\u001B[1;32m    203\u001B[0m \u001B[43m    \u001B[49m\u001B[43mattempts\u001B[49m\u001B[43m,\u001B[49m\n\u001B[1;32m    204\u001B[0m \u001B[43m    \u001B[49m\u001B[43moperation_model\u001B[49m\u001B[43m,\u001B[49m\n\u001B[1;32m    205\u001B[0m \u001B[43m    \u001B[49m\u001B[43mrequest_dict\u001B[49m\u001B[43m,\u001B[49m\n\u001B[1;32m    206\u001B[0m \u001B[43m    \u001B[49m\u001B[43msuccess_response\u001B[49m\u001B[43m,\u001B[49m\n\u001B[1;32m    207\u001B[0m \u001B[43m    \u001B[49m\u001B[43mexception\u001B[49m\u001B[43m,\u001B[49m\n\u001B[1;32m    208\u001B[0m \u001B[43m\u001B[49m\u001B[43m)\u001B[49m:\n\u001B[1;32m    209\u001B[0m     attempts \u001B[38;5;241m+\u001B[39m\u001B[38;5;241m=\u001B[39m \u001B[38;5;241m1\u001B[39m\n\u001B[1;32m    210\u001B[0m     \u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39m_update_retries_context(context, attempts, success_response)\n",
      "File \u001B[0;32m/opt/venv/lib/python3.9/site-packages/botocore/endpoint.py:354\u001B[0m, in \u001B[0;36mEndpoint._needs_retry\u001B[0;34m(self, attempts, operation_model, request_dict, response, caught_exception)\u001B[0m\n\u001B[1;32m    352\u001B[0m service_id \u001B[38;5;241m=\u001B[39m operation_model\u001B[38;5;241m.\u001B[39mservice_model\u001B[38;5;241m.\u001B[39mservice_id\u001B[38;5;241m.\u001B[39mhyphenize()\n\u001B[1;32m    353\u001B[0m event_name \u001B[38;5;241m=\u001B[39m \u001B[38;5;124mf\u001B[39m\u001B[38;5;124m\"\u001B[39m\u001B[38;5;124mneeds-retry.\u001B[39m\u001B[38;5;132;01m{\u001B[39;00mservice_id\u001B[38;5;132;01m}\u001B[39;00m\u001B[38;5;124m.\u001B[39m\u001B[38;5;132;01m{\u001B[39;00moperation_model\u001B[38;5;241m.\u001B[39mname\u001B[38;5;132;01m}\u001B[39;00m\u001B[38;5;124m\"\u001B[39m\n\u001B[0;32m--> 354\u001B[0m responses \u001B[38;5;241m=\u001B[39m \u001B[38;5;28;43mself\u001B[39;49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43m_event_emitter\u001B[49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43memit\u001B[49m\u001B[43m(\u001B[49m\n\u001B[1;32m    355\u001B[0m \u001B[43m    \u001B[49m\u001B[43mevent_name\u001B[49m\u001B[43m,\u001B[49m\n\u001B[1;32m    356\u001B[0m \u001B[43m    \u001B[49m\u001B[43mresponse\u001B[49m\u001B[38;5;241;43m=\u001B[39;49m\u001B[43mresponse\u001B[49m\u001B[43m,\u001B[49m\n\u001B[1;32m    357\u001B[0m \u001B[43m    \u001B[49m\u001B[43mendpoint\u001B[49m\u001B[38;5;241;43m=\u001B[39;49m\u001B[38;5;28;43mself\u001B[39;49m\u001B[43m,\u001B[49m\n\u001B[1;32m    358\u001B[0m \u001B[43m    \u001B[49m\u001B[43moperation\u001B[49m\u001B[38;5;241;43m=\u001B[39;49m\u001B[43moperation_model\u001B[49m\u001B[43m,\u001B[49m\n\u001B[1;32m    359\u001B[0m \u001B[43m    \u001B[49m\u001B[43mattempts\u001B[49m\u001B[38;5;241;43m=\u001B[39;49m\u001B[43mattempts\u001B[49m\u001B[43m,\u001B[49m\n\u001B[1;32m    360\u001B[0m \u001B[43m    \u001B[49m\u001B[43mcaught_exception\u001B[49m\u001B[38;5;241;43m=\u001B[39;49m\u001B[43mcaught_exception\u001B[49m\u001B[43m,\u001B[49m\n\u001B[1;32m    361\u001B[0m \u001B[43m    \u001B[49m\u001B[43mrequest_dict\u001B[49m\u001B[38;5;241;43m=\u001B[39;49m\u001B[43mrequest_dict\u001B[49m\u001B[43m,\u001B[49m\n\u001B[1;32m    362\u001B[0m \u001B[43m\u001B[49m\u001B[43m)\u001B[49m\n\u001B[1;32m    363\u001B[0m handler_response \u001B[38;5;241m=\u001B[39m first_non_none_response(responses)\n\u001B[1;32m    364\u001B[0m \u001B[38;5;28;01mif\u001B[39;00m handler_response \u001B[38;5;129;01mis\u001B[39;00m \u001B[38;5;28;01mNone\u001B[39;00m:\n",
      "File \u001B[0;32m/opt/venv/lib/python3.9/site-packages/botocore/hooks.py:412\u001B[0m, in \u001B[0;36mEventAliaser.emit\u001B[0;34m(self, event_name, **kwargs)\u001B[0m\n\u001B[1;32m    410\u001B[0m \u001B[38;5;28;01mdef\u001B[39;00m \u001B[38;5;21memit\u001B[39m(\u001B[38;5;28mself\u001B[39m, event_name, \u001B[38;5;241m*\u001B[39m\u001B[38;5;241m*\u001B[39mkwargs):\n\u001B[1;32m    411\u001B[0m     aliased_event_name \u001B[38;5;241m=\u001B[39m \u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39m_alias_event_name(event_name)\n\u001B[0;32m--> 412\u001B[0m     \u001B[38;5;28;01mreturn\u001B[39;00m \u001B[38;5;28;43mself\u001B[39;49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43m_emitter\u001B[49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43memit\u001B[49m\u001B[43m(\u001B[49m\u001B[43maliased_event_name\u001B[49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[38;5;241;43m*\u001B[39;49m\u001B[38;5;241;43m*\u001B[39;49m\u001B[43mkwargs\u001B[49m\u001B[43m)\u001B[49m\n",
      "File \u001B[0;32m/opt/venv/lib/python3.9/site-packages/botocore/hooks.py:256\u001B[0m, in \u001B[0;36mHierarchicalEmitter.emit\u001B[0;34m(self, event_name, **kwargs)\u001B[0m\n\u001B[1;32m    245\u001B[0m \u001B[38;5;28;01mdef\u001B[39;00m \u001B[38;5;21memit\u001B[39m(\u001B[38;5;28mself\u001B[39m, event_name, \u001B[38;5;241m*\u001B[39m\u001B[38;5;241m*\u001B[39mkwargs):\n\u001B[1;32m    246\u001B[0m     \u001B[38;5;124;03m\"\"\"\u001B[39;00m\n\u001B[1;32m    247\u001B[0m \u001B[38;5;124;03m    Emit an event by name with arguments passed as keyword args.\u001B[39;00m\n\u001B[1;32m    248\u001B[0m \n\u001B[0;32m   (...)\u001B[0m\n\u001B[1;32m    254\u001B[0m \u001B[38;5;124;03m             handlers.\u001B[39;00m\n\u001B[1;32m    255\u001B[0m \u001B[38;5;124;03m    \"\"\"\u001B[39;00m\n\u001B[0;32m--> 256\u001B[0m     \u001B[38;5;28;01mreturn\u001B[39;00m \u001B[38;5;28;43mself\u001B[39;49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43m_emit\u001B[49m\u001B[43m(\u001B[49m\u001B[43mevent_name\u001B[49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[43mkwargs\u001B[49m\u001B[43m)\u001B[49m\n",
      "File \u001B[0;32m/opt/venv/lib/python3.9/site-packages/botocore/hooks.py:239\u001B[0m, in \u001B[0;36mHierarchicalEmitter._emit\u001B[0;34m(self, event_name, kwargs, stop_on_response)\u001B[0m\n\u001B[1;32m    237\u001B[0m \u001B[38;5;28;01mfor\u001B[39;00m handler \u001B[38;5;129;01min\u001B[39;00m handlers_to_call:\n\u001B[1;32m    238\u001B[0m     logger\u001B[38;5;241m.\u001B[39mdebug(\u001B[38;5;124m'\u001B[39m\u001B[38;5;124mEvent \u001B[39m\u001B[38;5;132;01m%s\u001B[39;00m\u001B[38;5;124m: calling handler \u001B[39m\u001B[38;5;132;01m%s\u001B[39;00m\u001B[38;5;124m'\u001B[39m, event_name, handler)\n\u001B[0;32m--> 239\u001B[0m     response \u001B[38;5;241m=\u001B[39m \u001B[43mhandler\u001B[49m\u001B[43m(\u001B[49m\u001B[38;5;241;43m*\u001B[39;49m\u001B[38;5;241;43m*\u001B[39;49m\u001B[43mkwargs\u001B[49m\u001B[43m)\u001B[49m\n\u001B[1;32m    240\u001B[0m     responses\u001B[38;5;241m.\u001B[39mappend((handler, response))\n\u001B[1;32m    241\u001B[0m     \u001B[38;5;28;01mif\u001B[39;00m stop_on_response \u001B[38;5;129;01mand\u001B[39;00m response \u001B[38;5;129;01mis\u001B[39;00m \u001B[38;5;129;01mnot\u001B[39;00m \u001B[38;5;28;01mNone\u001B[39;00m:\n",
      "File \u001B[0;32m/opt/venv/lib/python3.9/site-packages/botocore/retryhandler.py:207\u001B[0m, in \u001B[0;36mRetryHandler.__call__\u001B[0;34m(self, attempts, response, caught_exception, **kwargs)\u001B[0m\n\u001B[1;32m    204\u001B[0m     retries_context \u001B[38;5;241m=\u001B[39m kwargs[\u001B[38;5;124m'\u001B[39m\u001B[38;5;124mrequest_dict\u001B[39m\u001B[38;5;124m'\u001B[39m][\u001B[38;5;124m'\u001B[39m\u001B[38;5;124mcontext\u001B[39m\u001B[38;5;124m'\u001B[39m]\u001B[38;5;241m.\u001B[39mget(\u001B[38;5;124m'\u001B[39m\u001B[38;5;124mretries\u001B[39m\u001B[38;5;124m'\u001B[39m)\n\u001B[1;32m    205\u001B[0m     checker_kwargs\u001B[38;5;241m.\u001B[39mupdate({\u001B[38;5;124m'\u001B[39m\u001B[38;5;124mretries_context\u001B[39m\u001B[38;5;124m'\u001B[39m: retries_context})\n\u001B[0;32m--> 207\u001B[0m \u001B[38;5;28;01mif\u001B[39;00m \u001B[38;5;28;43mself\u001B[39;49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43m_checker\u001B[49m\u001B[43m(\u001B[49m\u001B[38;5;241;43m*\u001B[39;49m\u001B[38;5;241;43m*\u001B[39;49m\u001B[43mchecker_kwargs\u001B[49m\u001B[43m)\u001B[49m:\n\u001B[1;32m    208\u001B[0m     result \u001B[38;5;241m=\u001B[39m \u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39m_action(attempts\u001B[38;5;241m=\u001B[39mattempts)\n\u001B[1;32m    209\u001B[0m     logger\u001B[38;5;241m.\u001B[39mdebug(\u001B[38;5;124m\"\u001B[39m\u001B[38;5;124mRetry needed, action of: \u001B[39m\u001B[38;5;132;01m%s\u001B[39;00m\u001B[38;5;124m\"\u001B[39m, result)\n",
      "File \u001B[0;32m/opt/venv/lib/python3.9/site-packages/botocore/retryhandler.py:284\u001B[0m, in \u001B[0;36mMaxAttemptsDecorator.__call__\u001B[0;34m(self, attempt_number, response, caught_exception, retries_context)\u001B[0m\n\u001B[1;32m    279\u001B[0m \u001B[38;5;28;01mif\u001B[39;00m retries_context:\n\u001B[1;32m    280\u001B[0m     retries_context[\u001B[38;5;124m'\u001B[39m\u001B[38;5;124mmax\u001B[39m\u001B[38;5;124m'\u001B[39m] \u001B[38;5;241m=\u001B[39m \u001B[38;5;28mmax\u001B[39m(\n\u001B[1;32m    281\u001B[0m         retries_context\u001B[38;5;241m.\u001B[39mget(\u001B[38;5;124m'\u001B[39m\u001B[38;5;124mmax\u001B[39m\u001B[38;5;124m'\u001B[39m, \u001B[38;5;241m0\u001B[39m), \u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39m_max_attempts\n\u001B[1;32m    282\u001B[0m     )\n\u001B[0;32m--> 284\u001B[0m should_retry \u001B[38;5;241m=\u001B[39m \u001B[38;5;28;43mself\u001B[39;49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43m_should_retry\u001B[49m\u001B[43m(\u001B[49m\n\u001B[1;32m    285\u001B[0m \u001B[43m    \u001B[49m\u001B[43mattempt_number\u001B[49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[43mresponse\u001B[49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[43mcaught_exception\u001B[49m\n\u001B[1;32m    286\u001B[0m \u001B[43m\u001B[49m\u001B[43m)\u001B[49m\n\u001B[1;32m    287\u001B[0m \u001B[38;5;28;01mif\u001B[39;00m should_retry:\n\u001B[1;32m    288\u001B[0m     \u001B[38;5;28;01mif\u001B[39;00m attempt_number \u001B[38;5;241m>\u001B[39m\u001B[38;5;241m=\u001B[39m \u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39m_max_attempts:\n\u001B[1;32m    289\u001B[0m         \u001B[38;5;66;03m# explicitly set MaxAttemptsReached\u001B[39;00m\n",
      "File \u001B[0;32m/opt/venv/lib/python3.9/site-packages/botocore/retryhandler.py:320\u001B[0m, in \u001B[0;36mMaxAttemptsDecorator._should_retry\u001B[0;34m(self, attempt_number, response, caught_exception)\u001B[0m\n\u001B[1;32m    316\u001B[0m         \u001B[38;5;28;01mreturn\u001B[39;00m \u001B[38;5;28;01mTrue\u001B[39;00m\n\u001B[1;32m    317\u001B[0m \u001B[38;5;28;01melse\u001B[39;00m:\n\u001B[1;32m    318\u001B[0m     \u001B[38;5;66;03m# If we've exceeded the max attempts we just let the exception\u001B[39;00m\n\u001B[1;32m    319\u001B[0m     \u001B[38;5;66;03m# propogate if one has occurred.\u001B[39;00m\n\u001B[0;32m--> 320\u001B[0m     \u001B[38;5;28;01mreturn\u001B[39;00m \u001B[38;5;28;43mself\u001B[39;49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43m_checker\u001B[49m\u001B[43m(\u001B[49m\u001B[43mattempt_number\u001B[49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[43mresponse\u001B[49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[43mcaught_exception\u001B[49m\u001B[43m)\u001B[49m\n",
      "File \u001B[0;32m/opt/venv/lib/python3.9/site-packages/botocore/retryhandler.py:363\u001B[0m, in \u001B[0;36mMultiChecker.__call__\u001B[0;34m(self, attempt_number, response, caught_exception)\u001B[0m\n\u001B[1;32m    361\u001B[0m \u001B[38;5;28;01mdef\u001B[39;00m \u001B[38;5;21m__call__\u001B[39m(\u001B[38;5;28mself\u001B[39m, attempt_number, response, caught_exception):\n\u001B[1;32m    362\u001B[0m     \u001B[38;5;28;01mfor\u001B[39;00m checker \u001B[38;5;129;01min\u001B[39;00m \u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39m_checkers:\n\u001B[0;32m--> 363\u001B[0m         checker_response \u001B[38;5;241m=\u001B[39m \u001B[43mchecker\u001B[49m\u001B[43m(\u001B[49m\n\u001B[1;32m    364\u001B[0m \u001B[43m            \u001B[49m\u001B[43mattempt_number\u001B[49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[43mresponse\u001B[49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[43mcaught_exception\u001B[49m\n\u001B[1;32m    365\u001B[0m \u001B[43m        \u001B[49m\u001B[43m)\u001B[49m\n\u001B[1;32m    366\u001B[0m         \u001B[38;5;28;01mif\u001B[39;00m checker_response:\n\u001B[1;32m    367\u001B[0m             \u001B[38;5;28;01mreturn\u001B[39;00m checker_response\n",
      "File \u001B[0;32m/opt/venv/lib/python3.9/site-packages/botocore/retryhandler.py:247\u001B[0m, in \u001B[0;36mBaseChecker.__call__\u001B[0;34m(self, attempt_number, response, caught_exception)\u001B[0m\n\u001B[1;32m    245\u001B[0m     \u001B[38;5;28;01mreturn\u001B[39;00m \u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39m_check_response(attempt_number, response)\n\u001B[1;32m    246\u001B[0m \u001B[38;5;28;01melif\u001B[39;00m caught_exception \u001B[38;5;129;01mis\u001B[39;00m \u001B[38;5;129;01mnot\u001B[39;00m \u001B[38;5;28;01mNone\u001B[39;00m:\n\u001B[0;32m--> 247\u001B[0m     \u001B[38;5;28;01mreturn\u001B[39;00m \u001B[38;5;28;43mself\u001B[39;49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43m_check_caught_exception\u001B[49m\u001B[43m(\u001B[49m\n\u001B[1;32m    248\u001B[0m \u001B[43m        \u001B[49m\u001B[43mattempt_number\u001B[49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[43mcaught_exception\u001B[49m\n\u001B[1;32m    249\u001B[0m \u001B[43m    \u001B[49m\u001B[43m)\u001B[49m\n\u001B[1;32m    250\u001B[0m \u001B[38;5;28;01melse\u001B[39;00m:\n\u001B[1;32m    251\u001B[0m     \u001B[38;5;28;01mraise\u001B[39;00m \u001B[38;5;167;01mValueError\u001B[39;00m(\u001B[38;5;124m\"\u001B[39m\u001B[38;5;124mBoth response and caught_exception are None.\u001B[39m\u001B[38;5;124m\"\u001B[39m)\n",
      "File \u001B[0;32m/opt/venv/lib/python3.9/site-packages/botocore/retryhandler.py:416\u001B[0m, in \u001B[0;36mExceptionRaiser._check_caught_exception\u001B[0;34m(self, attempt_number, caught_exception)\u001B[0m\n\u001B[1;32m    408\u001B[0m \u001B[38;5;28;01mdef\u001B[39;00m \u001B[38;5;21m_check_caught_exception\u001B[39m(\u001B[38;5;28mself\u001B[39m, attempt_number, caught_exception):\n\u001B[1;32m    409\u001B[0m     \u001B[38;5;66;03m# This is implementation specific, but this class is useful by\u001B[39;00m\n\u001B[1;32m    410\u001B[0m     \u001B[38;5;66;03m# coordinating with the MaxAttemptsDecorator.\u001B[39;00m\n\u001B[0;32m   (...)\u001B[0m\n\u001B[1;32m    414\u001B[0m     \u001B[38;5;66;03m# the MaxAttemptsDecorator is not interested in retrying the exception\u001B[39;00m\n\u001B[1;32m    415\u001B[0m     \u001B[38;5;66;03m# then this exception just propogates out past the retry code.\u001B[39;00m\n\u001B[0;32m--> 416\u001B[0m     \u001B[38;5;28;01mraise\u001B[39;00m caught_exception\n",
      "File \u001B[0;32m/opt/venv/lib/python3.9/site-packages/botocore/endpoint.py:281\u001B[0m, in \u001B[0;36mEndpoint._do_get_response\u001B[0;34m(self, request, operation_model, context)\u001B[0m\n\u001B[1;32m    279\u001B[0m     http_response \u001B[38;5;241m=\u001B[39m first_non_none_response(responses)\n\u001B[1;32m    280\u001B[0m     \u001B[38;5;28;01mif\u001B[39;00m http_response \u001B[38;5;129;01mis\u001B[39;00m \u001B[38;5;28;01mNone\u001B[39;00m:\n\u001B[0;32m--> 281\u001B[0m         http_response \u001B[38;5;241m=\u001B[39m \u001B[38;5;28;43mself\u001B[39;49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43m_send\u001B[49m\u001B[43m(\u001B[49m\u001B[43mrequest\u001B[49m\u001B[43m)\u001B[49m\n\u001B[1;32m    282\u001B[0m \u001B[38;5;28;01mexcept\u001B[39;00m HTTPClientError \u001B[38;5;28;01mas\u001B[39;00m e:\n\u001B[1;32m    283\u001B[0m     \u001B[38;5;28;01mreturn\u001B[39;00m (\u001B[38;5;28;01mNone\u001B[39;00m, e)\n",
      "File \u001B[0;32m/opt/venv/lib/python3.9/site-packages/botocore/endpoint.py:377\u001B[0m, in \u001B[0;36mEndpoint._send\u001B[0;34m(self, request)\u001B[0m\n\u001B[1;32m    376\u001B[0m \u001B[38;5;28;01mdef\u001B[39;00m \u001B[38;5;21m_send\u001B[39m(\u001B[38;5;28mself\u001B[39m, request):\n\u001B[0;32m--> 377\u001B[0m     \u001B[38;5;28;01mreturn\u001B[39;00m \u001B[38;5;28;43mself\u001B[39;49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43mhttp_session\u001B[49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43msend\u001B[49m\u001B[43m(\u001B[49m\u001B[43mrequest\u001B[49m\u001B[43m)\u001B[49m\n",
      "File \u001B[0;32m/opt/venv/lib/python3.9/site-packages/botocore/httpsession.py:484\u001B[0m, in \u001B[0;36mURLLib3Session.send\u001B[0;34m(self, request)\u001B[0m\n\u001B[1;32m    482\u001B[0m     \u001B[38;5;28;01mraise\u001B[39;00m SSLError(endpoint_url\u001B[38;5;241m=\u001B[39mrequest\u001B[38;5;241m.\u001B[39murl, error\u001B[38;5;241m=\u001B[39me)\n\u001B[1;32m    483\u001B[0m \u001B[38;5;28;01mexcept\u001B[39;00m (NewConnectionError, socket\u001B[38;5;241m.\u001B[39mgaierror) \u001B[38;5;28;01mas\u001B[39;00m e:\n\u001B[0;32m--> 484\u001B[0m     \u001B[38;5;28;01mraise\u001B[39;00m EndpointConnectionError(endpoint_url\u001B[38;5;241m=\u001B[39mrequest\u001B[38;5;241m.\u001B[39murl, error\u001B[38;5;241m=\u001B[39me)\n\u001B[1;32m    485\u001B[0m \u001B[38;5;28;01mexcept\u001B[39;00m ProxyError \u001B[38;5;28;01mas\u001B[39;00m e:\n\u001B[1;32m    486\u001B[0m     \u001B[38;5;28;01mraise\u001B[39;00m ProxyConnectionError(\n\u001B[1;32m    487\u001B[0m         proxy_url\u001B[38;5;241m=\u001B[39mmask_proxy_url(proxy_url), error\u001B[38;5;241m=\u001B[39me\n\u001B[1;32m    488\u001B[0m     )\n",
      "\u001B[0;31mEndpointConnectionError\u001B[0m: Could not connect to the endpoint URL: \"https://avalon-benchmark.s3.us-west-2.amazonaws.com/random_policy__0__avalon_worlds__0824_full.tar.gz\""
     ]
    }
   ],
   "source": [
    "\n",
    "## RANDOM BASELINE\n",
    "\n",
    "random_result_keys = [f\"random_policy__{i}__{FIXED_WORLD_KEY}\" for i in range(10)]\n",
    "combined_random_results = defaultdict(list)\n",
    "for result_key in random_result_keys:\n",
    "    reported_data = json.loads(s3_client.load(result_key))\n",
    "    for k, v in reported_data[\"all_results\"].items():\n",
    "        combined_random_results[k].append(v)\n",
    "\n",
    "random_agg_score = {}\n",
    "for k, v in combined_random_results.items():\n",
    "    random_agg_score[k] = np.mean(v)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "\n",
    "## HUMAN BASELINE\n",
    "human_scores_key = \"avalon_human_scores__0826\"\n",
    "\n",
    "# to upload:\n",
    "# with open(\"/mnt/private/avalon/human_scores__0824.json\", \"r\") as f:\n",
    "#     human_scores = json.load(f)\n",
    "# s3_client.save(human_scores_key, json.dumps(human_scores))\n",
    "\n",
    "human_scores = json.loads(s3_client.load(human_scores_key))\n",
    "human_agg_score = {}\n",
    "folder_name_lookup = {}\n",
    "for world_dir, human_results_on_world in human_scores.items():\n",
    "    k = str(int(world_dir.split(\"__\")[1]))\n",
    "    folder_name_lookup[k] = world_dir\n",
    "    human_agg_score[k] = np.mean(list(human_results_on_world.values()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "\n",
    "manually_excluded = set()\n",
    "valid_levels = set(human_agg_score.keys())  # & set(random_agg_score.keys())\n",
    "human_fail_levels = set()\n",
    "for k in valid_levels:\n",
    "    if human_agg_score[k] <= random_agg_score[k]:\n",
    "        human_fail_levels.add(k)\n",
    "        fail_folder = folder_name_lookup[k]\n",
    "        logger.warning(f\"No human succeeded at level {fail_folder}! {human_scores[fail_folder]}\")\n",
    "\n",
    "valid_levels -= human_fail_levels\n",
    "valid_levels -= manually_excluded\n",
    "\n",
    "\n",
    "def get_human_normalized_score(world_index: str, raw_score: float) -> float:\n",
    "    if world_index not in valid_levels:\n",
    "        logger.warning(f\"{world_index} not in valid_levels returning unnormalized score!\")\n",
    "        return raw_score\n",
    "    return max(0, raw_score - random_agg_score[world_index]) / (\n",
    "        human_agg_score[world_index] - random_agg_score[world_index]\n",
    "    )\n",
    "\n",
    "\n",
    "world_id_by_task = defaultdict(list)\n",
    "simple_tasks = set(x.value.lower() for x in TASKS_BY_TASK_GROUP[AvalonTaskGroup.SIMPLE])\n",
    "compositional_tasks = set(x.value.lower() for x in TASKS_BY_TASK_GROUP[AvalonTaskGroup.COMPOSITIONAL])\n",
    "for key in valid_levels:\n",
    "    task = folder_name_lookup[key].split(\"__\")[0]\n",
    "    world_id_by_task[task].append(key)\n",
    "    if task in simple_tasks:\n",
    "        world_id_by_task[\"simple\"].append(key)\n",
    "    if task in compositional_tasks:\n",
    "        world_id_by_task[\"compositional\"].append(key)\n",
    "    world_id_by_task[\"all\"].append(key)\n",
    "\n",
    "\n",
    "def task_dict_to_score_numpys(task_dict: Dict[str, List[float]]):\n",
    "    np_task_dict: Dict[str, np.ndarray] = {}\n",
    "    for task, keys in world_id_by_task.items():\n",
    "        np_task_dict[task] = np.stack([np.array(task_dict[k]) for k in keys])\n",
    "    return np_task_dict\n",
    "\n",
    "\n",
    "PRINT_TASK_ORDER = [x.value.lower() for x in TASKS_BY_TASK_GROUP[AvalonTaskGroup.ALL]] + [\n",
    "    \"simple\",\n",
    "    \"compositional\",\n",
    "    \"all\",\n",
    "]\n",
    "\n",
    "\n",
    "aggregate_func = lambda x: np.array(\n",
    "    [\n",
    "        metrics.aggregate_mean(x),\n",
    "        metrics.aggregate_optimality_gap(x),\n",
    "    ]\n",
    ")\n",
    "\n",
    "\n",
    "def print_aggregate_results(score_numpys: Dict[str, np.ndarray]):\n",
    "    aggregate_scores, aggregate_score_cis = rly.get_interval_estimates(score_numpys, aggregate_func, reps=1000)\n",
    "    for task in PRINT_TASK_ORDER:\n",
    "        if task not in aggregate_scores:\n",
    "            continue\n",
    "        mean, opt_gap = aggregate_scores[task]\n",
    "        mean_low, mean_high = aggregate_score_cis[task][:, 0]\n",
    "        mean_sig = (mean_high - mean_low) / 2\n",
    "        opt_gap_low, opt_gap_high = aggregate_score_cis[task][:, 1]\n",
    "        opt_gap_sig = (mean_high - mean_low) / 2\n",
    "        logger.info(f\"{task.upper()}: mean {mean:.3f} +- {mean_sig:.3f}, og: {opt_gap:.3f} +- {opt_gap_sig:.3f}\")\n",
    "\n",
    "\n",
    "def latex_aggregate_results(score_numpys: Dict[str, np.ndarray]) -> Dict[str, Dict[str, str]]:\n",
    "    lines = {}\n",
    "    aggregate_scores, aggregate_score_cis = rly.get_interval_estimates(score_numpys, aggregate_func, reps=1000)\n",
    "    for task in PRINT_TASK_ORDER:\n",
    "        if task not in aggregate_scores:\n",
    "            lines[task] = {\n",
    "                \"mean\": \"-\",\n",
    "                \"optimality_gap\": \"-\",\n",
    "            }\n",
    "            continue\n",
    "        mean, opt_gap = aggregate_scores[task]\n",
    "        mean_low, mean_high = aggregate_score_cis[task][:, 0]\n",
    "        mean_sig = (mean_high - mean_low) / 2\n",
    "        opt_gap_low, opt_gap_high = aggregate_score_cis[task][:, 1]\n",
    "        opt_gap_sig = (mean_high - mean_low) / 2\n",
    "        lines[task] = {\n",
    "            \"mean\": f\"{mean:.3f} $\\pm$ {mean_sig:.3f}\",\n",
    "            \"optimality_gap\": f\"{opt_gap:.3f} $\\pm$ {opt_gap_sig:.3f}\",\n",
    "        }\n",
    "    return lines"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "\n",
    "\n",
    "FINISHED_RUNS = {\n",
    "    \"single_task_navigate\": [\n",
    "        \"d2e999ac-539c-480e-9286-843e52152f93\",\n",
    "        \"d141c933-1120-4a65-b759-2e2bcf46240d\",\n",
    "        \"98802cc8-8a41-41b9-9a93-2a8545bca7cb\",\n",
    "        \"7a37d446-5391-47c9-9742-2fd679b17b1a\",\n",
    "        \"46b9dd6e-5d6e-4b1d-a500-cc70b2afac02\",\n",
    "        \"5b1c2b4a-d43a-4680-bba7-a7be4f2d8915\",\n",
    "    ],\n",
    "    \"single_task_carry\": [\n",
    "        \"8222027c-49b2-4950-aa14-6a4d612f2f70\",\n",
    "        \"95fe5e36-729e-462e-9902-5305db492d15\",\n",
    "        \"f04bda8e-e940-4edc-9e70-bbfd7e1d7253\",\n",
    "        \"d78e3ae2-c978-42ac-b7d6-0a3c75287452\",\n",
    "        \"9970abc0-9a50-46d7-bd42-5ab7fe8a23c0\",\n",
    "        \"1ea2118a-7418-4c99-8c25-919e3eeb20a2\",\n",
    "    ],\n",
    "    \"single_task_bridge\": [\n",
    "        \"02e3b730-7f87-4f93-be01-d77d12adba04\",\n",
    "        \"42dac847-fc97-4f6e-a111-6cd67b06ff0e\",\n",
    "        \"f7e94e54-43be-4f7a-bd49-d8d8ccd3c15f\",\n",
    "        \"b506fb63-897e-491c-ad3f-0199cec3fc50\",\n",
    "        \"3cb5f25b-1b1c-4e6d-9d34-c17966954af3\",\n",
    "    ],\n",
    "    \"single_task_stack\": [\n",
    "        \"4f63e2d0-4a17-4911-a819-20ea2c77c8f4\",\n",
    "        \"d56aea61-4077-4574-a6ff-c75340825d10\",\n",
    "        \"008819b4-fc99-41e8-80c2-1a3571875053\",\n",
    "        \"b3b92641-e981-4133-89b3-0cb075c7b141\",\n",
    "        \"3e587d03-18ae-4131-ba0a-e680cbb0924b\",\n",
    "    ],\n",
    "    \"single_task_open\": [\n",
    "        \"2d748ad1-f250-455a-b5a8-5ddfbf1537d7\",\n",
    "        \"f730db29-fadb-4d04-9ebf-62f8e4e5e271\",\n",
    "        \"1e771c08-6a2a-4ed8-abe7-c372a3875710\",\n",
    "        \"131fb969-7262-4f36-860d-61122016eefa\",\n",
    "        \"18f30685-5c54-4bc1-903f-844bac12187d\",\n",
    "        \"d4399557-1e4d-4452-9436-11a8c2fabc05\",\n",
    "    ],\n",
    "    \"single_task_push\": [\n",
    "        \"528b6372-9470-4859-8f52-d66a3b1ec986\",\n",
    "        \"15a9f45a-319e-4254-a892-975c98bfbc95\",\n",
    "        \"d462b56e-5228-427d-93c3-e347fd8154f5\",\n",
    "        \"8e8870cd-cfa3-483e-b93d-fa089ab7e51b\",\n",
    "        \"5f19829d-f0e6-4c25-b5b2-03d6eb1f701d\",\n",
    "    ],\n",
    "    \"single_task_fight\": [\n",
    "        \"f3ff80c8-fbc3-4277-bdde-364ded17407f\",\n",
    "        \"26edb4a1-d4f0-47c8-97f5-f3bba1344f67\",\n",
    "        \"f19628c2-2084-41dc-8618-952a137f28df\",\n",
    "        \"da2aa3aa-6361-4c32-99f5-926c6891cea9\",\n",
    "        \"b5599519-8e81-479b-a670-1c9059f2db1f\",\n",
    "    ],\n",
    "    \"single_task_hunt\": [\n",
    "        \"fb44000f-b1ec-48a8-8c60-9a332cb64f06\",\n",
    "        \"b7a0e8a9-4caa-4b4d-8339-3c5c3d012893\",\n",
    "        \"42144a8b-baa4-44f0-8752-0aeecaaeb868\",\n",
    "        \"990bc5b7-5895-4c70-8ef8-70ce28a82a1d\",\n",
    "        \"b96c9bbe-ade4-4ea0-bd96-2437b77daa24\",\n",
    "        \"e4b87d99-7187-44bd-a636-d7900705b71b\",\n",
    "    ],\n",
    "    \"single_task_avoid\": [\n",
    "        \"05e7cccc-cfc9-4345-a5f2-ea43f534a289\",\n",
    "        \"8e65b732-2d04-4075-a714-9442af3e736b\",\n",
    "        \"e8b864c9-6ac4-4a7d-99b2-fa85cc2b3010\",\n",
    "        \"9b7ebfb4-6cf4-4fa9-a373-3cbc118389b0\",\n",
    "        \"edf5e5d7-2fd0-43a5-a1cc-0276287af992\",\n",
    "        \"556771f1-0579-4a64-a8be-ffcc1bd2b1c2\",\n",
    "    ],\n",
    "    \"single_task_throw\": [\n",
    "        \"79c194f7-569d-40dd-9c82-1703ff799cfd\",\n",
    "        \"6954376e-b840-4053-98ac-89336c4174bb\",\n",
    "        \"9ec8cc8f-d318-45ac-ba8d-cc7d6c1b801c\",\n",
    "        \"0bf62354-33e2-48df-96a7-2c2a1f73fedf\",\n",
    "        \"e9ef062a-cd64-445e-b086-b15c6f33b346\",\n",
    "        \"18e1eb11-e44d-457e-a7a8-2522ede2871e\",\n",
    "    ],\n",
    "    \"single_task_descend\": [\n",
    "        \"a663d1a4-0e03-423a-a8db-1a5675b7d0d0\",\n",
    "        \"1bce0f94-9924-4c30-8e0f-ee966e14ea34\",\n",
    "        \"94e63f45-a7fb-4dce-beb2-01a97af2e0e6\",\n",
    "        \"0cc8a42a-e73e-4c88-bd94-6a2d4eaa46f5\",\n",
    "        \"0d915730-3795-4703-8834-01b74ff37f72\",\n",
    "        \"9cfeb379-a59f-42cd-93df-2368fe96903b\",\n",
    "    ],\n",
    "    \"no_curr\": [\n",
    "        \"2182d660-339d-4b59-a78a-f36b8be591d9\",\n",
    "        \"9756bca1-a589-4644-a416-cc234904bc78\",\n",
    "        \"3a610c73-5d22-4582-a9f8-6444dfcc4feb\",\n",
    "        \"a18738ad-0fc4-4629-8e9e-5d6170dadd56\",\n",
    "        \"a4932faa-638e-420e-bf2a-0f79b2cc52de\",\n",
    "        \"f90b498a-15fe-40e8-8606-66ab8d10a445\",\n",
    "    ],\n",
    "    \"single_task_climb\": [\n",
    "        \"076bfe0d-48ae-4033-ae28-c796a28bda51\",\n",
    "        \"f13797b0-fa33-4d52-a8ab-170e9ceb1950\",\n",
    "        \"47c7b009-f583-45b3-b104-dcb64e427129\",\n",
    "        \"2c38e712-d32c-4a1f-9cbb-02560f911c49\",\n",
    "        \"d08f208e-b237-4d0f-a38b-8dcd589226df\",\n",
    "        \"e0097c76-332a-463c-8e17-718040eed51c\",\n",
    "    ],\n",
    "    \"single_task_scramble\": [\n",
    "        \"39f106fe-a507-4545-9016-a72173328da2\",\n",
    "        \"333f76e4-8c9a-4f82-bf73-5a7992b7bba7\",\n",
    "        \"82b39ad2-3c14-4151-b8d0-a464640db8e8\",\n",
    "        \"598acb50-4674-4dfb-8fb6-3c5ec2be5d75\",\n",
    "        \"e0f267ce-4021-458e-b41d-0b81da1790bb\",\n",
    "        \"cf9a1fad-9cd9-4e6e-b462-1955e17a84d8\",\n",
    "    ],\n",
    "    \"single_task_explore\": [\n",
    "        \"eeb7ecbb-5303-4922-a236-35a86c8aa2eb\",\n",
    "        \"9ebb3ac6-40e2-4a79-8282-a6e9f4e7ec9d\",\n",
    "        # \"da1b4473-b78c-4891-bb57-7f52d81ae395\",\n",
    "        \"d042966c-f793-487b-9a2e-35e121520ddb\",\n",
    "        \"00bc5be8-6fbc-4802-821d-f5c6bb71119b\",\n",
    "        \"3c300c1b-2f7b-4e2e-b959-56655ec31ea1\",\n",
    "    ],\n",
    "    \"single_task_jump\": [\n",
    "        \"cf6ed141-ba00-4b8d-ae9c-8fe6217ed937\",\n",
    "        \"e15ba89e-fd78-465a-b47f-9d76226ffb5c\",\n",
    "        \"cf1efce4-f2f4-46f1-a2fa-c032868543e5\",\n",
    "        \"19c8d3ab-7e08-485b-8950-b323b8f0f4d6\",\n",
    "        \"168f5d1e-cb76-40f7-97a9-e2866c475038\",\n",
    "        \"b477163c-32c6-4d7b-9551-dbe626fba751\",\n",
    "    ],\n",
    "    \"single_task_move\": [\n",
    "        \"27d04e34-db0c-4d03-a5a0-1aa0d75df6a7\",\n",
    "        \"067d3abd-0388-4421-98c7-419014b8360d\",\n",
    "        \"45f62ceb-e0ff-4cc7-97a6-c75bdc0fa44e\",\n",
    "        \"be7f951b-6645-4b38-8ab9-b7c9fc207ff8\",\n",
    "        \"717909ba-9b05-4ade-8358-b6686efad897\",\n",
    "        \"e8d2bf33-ec8d-4a1b-a6c3-63342330507e\",\n",
    "    ],\n",
    "    \"single_task_eat\": [\n",
    "        \"295908c0-749c-400b-b207-2bdf335f0d13\",\n",
    "        \"461591f5-6da3-4314-a275-76f3806074a6\",\n",
    "        \"2eccc6b2-c348-4c5a-b585-d87ed7c99e85\",\n",
    "        \"4771fb18-5a48-415a-933a-d76f60ca9dd9\",\n",
    "        \"91729854-adfc-4630-96ff-60a7eb1d0f0c\",\n",
    "        \"ce1ca135-59f5-446c-911d-1047113e903f\",\n",
    "    ],\n",
    "    \"multi_task_compositional\": [\n",
    "        \"299772bd-5607-42d4-9b54-c497b420dfa4\",\n",
    "        \"7be73111-92a7-4a5d-8b06-4dd1381c6823\",\n",
    "        \"0ef0fb63-cbd3-460c-be59-a724c4c54ee1\",\n",
    "        \"76afa191-45f1-43dc-a821-717ce215a1fe\",\n",
    "        \"5f0dacb8-9067-4192-80a6-e34ff047dcd8\",\n",
    "        \"8050a696-2ee8-485e-a1f6-206f684baf8e\",\n",
    "    ],\n",
    "    \"multi_task_basic\": [\n",
    "        \"00b11b8a-7e74-4ebd-87e6-82ae500ca2ab\",\n",
    "        \"f99fecba-9827-4b10-9569-bde0f98b7cca\",\n",
    "        \"766165dd-57a5-4db4-a1ed-109d7229a832\",\n",
    "        \"6d6a7a81-36a4-4916-9e21-997a843dbd0c\",\n",
    "        \"70dbe576-71f8-494b-a820-b3aafd0480c5\",\n",
    "        \"f4c7614f-c973-4c3b-baa8-7a7c532fbc74\",\n",
    "    ],\n",
    "    \"multi_task_all\": [\n",
    "        \"9c225872-d04d-420c-b94e-1df06c49aa62\",\n",
    "        \"d0bfd187-28ac-4dc5-ae83-2af3ef605aef\",\n",
    "        \"ec2ce66c-be85-4ebc-b69d-55569c72e168\",\n",
    "        \"c896839d-abc7-4b90-8a5f-3104ea554842\",\n",
    "        \"65e7233f-08d2-4816-8549-92eaf58d8337\",\n",
    "        \"f5a38e34-3ba4-4e3d-9615-b920e8f087d8\",\n",
    "    ],\n",
    "    \"single_task_survive\": [\n",
    "        \"a3e47df0-eee7-499a-823a-3cada24aa2d3\",\n",
    "        \"a5e8f34f-61dc-4db8-b811-f2e8617d60d2\",\n",
    "        \"e1722229-9ceb-4ef7-9a2d-8bbbf81f0b4c\",\n",
    "        \"f31efa63-00e5-4a6c-803e-9ebd09096801\",\n",
    "        \"7212c774-e38b-4f92-91c3-039b0573f9ac\",\n",
    "    ],\n",
    "    \"single_task_gather\": [\n",
    "        \"3a741d47-3aed-4a7c-89b7-c48a637e0cc8\",\n",
    "        \"51ab7b25-926c-4153-9393-cce90a41a26c\",\n",
    "        \"fc8fd934-dd60-4f34-9b38-fb5010934422\",\n",
    "        \"d6dca059-eecf-499a-91dd-12ce6633449b\",\n",
    "        \"a1e3c387-4fda-4e11-a7bb-b8dec1cc7cf0\",\n",
    "    ],\n",
    "    \"single_task_find\": [\n",
    "        \"9695f3fa-4439-419a-b27f-48e45c4c7bcf\",\n",
    "        \"5d42fe7d-678c-4ad7-a2a6-c814c45f2783\",\n",
    "        \"c4bf76e8-b50e-45f9-9036-ceb7859d568c\",\n",
    "        \"619aa9cd-da54-4e93-8904-963b4ef21a67\",\n",
    "        \"81119314-2e65-411c-9b29-ec586c774574\",\n",
    "    ],\n",
    "    \"500m\": [\n",
    "        \"42c4f566-49d1-45c6-a4f2-1185ac49238c\",\n",
    "        \"cdb06a43-4dfb-4d52-878d-8495075a4a38\",\n",
    "        \"bd10b76d-ce5b-4e0a-8514-47d75e777ae6\",\n",
    "    ],\n",
    "}\n",
    "\n",
    "\n",
    "def get_wandb_result_key(xpid: str, fixed_world_key: str, checkpoint: str = \"final\"):\n",
    "    return f\"avalon_eval__{xpid}__{fixed_world_key}__{checkpoint}\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "\n",
    "## TORCHBEAST BASELINE\n",
    "\n",
    "wandb_runs = FINISHED_RUNS[\"multi_task_basic\"]\n",
    "combined_torchbeast_results = defaultdict(list)\n",
    "raw_torchbeast_datas = []\n",
    "for uuid in wandb_runs:\n",
    "    result_key = get_wandb_result_key(uuid, FIXED_WORLD_KEY)\n",
    "    reported_data = json.loads(s3_client.load(result_key))\n",
    "    raw_torchbeast_datas.append(reported_data)\n",
    "    missing_levels = valid_levels - set(reported_data[\"all_results\"].keys())\n",
    "    if len(missing_levels) > 0:\n",
    "        logger.warning(f\"{uuid} is missing {len(missing_levels)} levels, ignoring incomplete data\")\n",
    "    else:\n",
    "        for k in valid_levels:\n",
    "            level_data = reported_data[\"all_results\"][k]\n",
    "            combined_torchbeast_results[k].append(get_human_normalized_score(k, level_data))\n",
    "\n",
    "torchbeast_score_numpys = task_dict_to_score_numpys(combined_torchbeast_results)\n",
    "\n",
    "logger.info(\"\\nTorchbeast 50m aggregated scores\")\n",
    "print_aggregate_results(torchbeast_score_numpys)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "\n",
    "\n",
    "## PPO BASELINE\n",
    "result_keys = [\n",
    "    \"avalon_eval__zack_zack_ppo_3t0oz4qc_final.pt__avalon_worlds__0824_full.tar.gz__final\",\n",
    "    \"avalon_eval__zack_zack_ppo_3ox73b4a_final.pt__avalon_worlds__0824_full.tar.gz__final\",\n",
    "    \"avalon_eval__zack_zack_ppo_fhgppehq_final.pt__avalon_worlds__0824_full.tar.gz__final\",\n",
    "    \"avalon_eval__zack_zack_ppo_1acn5jsy_final.pt__avalon_worlds__0824_full.tar.gz__final\",\n",
    "    \"avalon_eval__zack_zack_ppo_1wli6rk8_final.pt__avalon_worlds__0824_full.tar.gz__final\",\n",
    "]\n",
    "combined_ppo_results = defaultdict(list)\n",
    "raw_ppo_datas = []\n",
    "for result_key in result_keys:\n",
    "    reported_data = json.loads(s3_client.load(result_key))\n",
    "    raw_ppo_datas.append(reported_data)\n",
    "    missing_levels = valid_levels - set(reported_data[\"all_results\"].keys())\n",
    "    if len(missing_levels) > 0:\n",
    "        logger.warning(f\"{result_key} is missing {len(missing_levels)} levels, ignoring incomplete data\")\n",
    "    else:\n",
    "        for k in valid_levels:\n",
    "            level_data = reported_data[\"all_results\"][k]\n",
    "            combined_ppo_results[k].append(get_human_normalized_score(k, level_data))\n",
    "\n",
    "\n",
    "ppo_score_numpys = task_dict_to_score_numpys(combined_ppo_results)\n",
    "logger.info(\"\\nPPO 50m aggregated scores\")\n",
    "print_aggregate_results(ppo_score_numpys)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "\n",
    "\n",
    "## Dreamer BASELINE\n",
    "result_keys = [\n",
    "    \"avalon_eval__zack_dreamer_new_z2r75675_final.pt__avalon_worlds__0824_full.tar.gz__final\",\n",
    "    \"avalon_eval__zack_dreamer_new_385c6vbz_final.pt__avalon_worlds__0824_full.tar.gz__final\",\n",
    "    \"avalon_eval__zack_dreamer_new_lb9b0hut_final.pt__avalon_worlds__0824_full.tar.gz__final\",\n",
    "    \"avalon_eval__zack_dreamer_new_9o64zy4y_final.pt__avalon_worlds__0824_full.tar.gz__final\",\n",
    "    \"avalon_eval__zack_dreamer_new_29rrbilj_final.pt__avalon_worlds__0824_full.tar.gz__final\",\n",
    "]\n",
    "combined_dreamer_results = defaultdict(list)\n",
    "raw_dreamer_datas = []\n",
    "for result_key in result_keys:\n",
    "    reported_data = json.loads(s3_client.load(result_key))\n",
    "    raw_dreamer_datas.append(reported_data)\n",
    "    missing_levels = valid_levels - set(reported_data[\"all_results\"].keys())\n",
    "    if len(missing_levels) > 0:\n",
    "        logger.warning(f\"{result_key} is missing {len(missing_levels)} levels, ignoring incomplete data\")\n",
    "    else:\n",
    "        for k in valid_levels:\n",
    "            level_data = reported_data[\"all_results\"][k]\n",
    "            combined_dreamer_results[k].append(get_human_normalized_score(k, level_data))\n",
    "\n",
    "\n",
    "dreamer_score_numpys = task_dict_to_score_numpys(combined_dreamer_results)\n",
    "logger.info(\"\\nDreamer 50m aggregated scores\")\n",
    "print_aggregate_results(dreamer_score_numpys)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "\n",
    "## TORCHBEAST CURRICULUM ABLATION\n",
    "\n",
    "torchbeast_curriculum_results = {\n",
    "    \"500m\": defaultdict(list),\n",
    "    \"no_curr\": defaultdict(list),\n",
    "    \"multi_task_compositional\": defaultdict(list),\n",
    "    \"multi_task_all\": defaultdict(list),\n",
    "}\n",
    "for curriculum in torchbeast_curriculum_results.keys():\n",
    "    for uuid in FINISHED_RUNS[curriculum]:\n",
    "        result_key = get_wandb_result_key(uuid, FIXED_WORLD_KEY)\n",
    "        reported_data = json.loads(s3_client.load(result_key))\n",
    "        missing_levels = valid_levels - set(reported_data[\"all_results\"].keys())\n",
    "        if len(missing_levels) > 0:\n",
    "            logger.warning(f\"{result_key} is missing {len(missing_levels)} levels, ignoring incomplete data\")\n",
    "        else:\n",
    "            for k in valid_levels:\n",
    "                level_data = reported_data[\"all_results\"][k]\n",
    "                torchbeast_curriculum_results[curriculum][k].append(get_human_normalized_score(k, level_data))\n",
    "\n",
    "\n",
    "curriciulum_score_numpys = {k: task_dict_to_score_numpys(v) for k, v in torchbeast_curriculum_results.items()}\n",
    "logger.info(\"\\nCurriculum ablation aggregated scores\")\n",
    "for curriculum, scores in curriciulum_score_numpys.items():\n",
    "    logger.info(f'\\nCurriculum: {curriculum.replace(\"_\", \" \").capitalize()}')\n",
    "    print_aggregate_results(scores)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "\n",
    "## TORCHBEAST SINGLE TASK ABLATIONS\n",
    "\n",
    "torchbeast_single_task_results = {\n",
    "    k.split(\"_\")[-1]: defaultdict(list) for k in FINISHED_RUNS.keys() if k.startswith(\"single_task\")\n",
    "}\n",
    "\n",
    "for task in torchbeast_single_task_results.keys():\n",
    "    for uuid in FINISHED_RUNS[f\"single_task_{task}\"]:\n",
    "        result_key = get_wandb_result_key(uuid, FIXED_WORLD_KEY)\n",
    "        try:\n",
    "            reported_data = json.loads(s3_client.load(result_key))\n",
    "        except KeyError:\n",
    "            logger.warning(f\"Missing result key {result_key}\")\n",
    "            continue\n",
    "        missing_levels = valid_levels - set(reported_data[\"all_results\"].keys())\n",
    "        if len(missing_levels) > 0:\n",
    "            logger.warning(f\"{result_key} is missing {len(missing_levels)} levels, ignoring incomplete data\")\n",
    "        else:\n",
    "            for k in valid_levels:\n",
    "                level_data = reported_data[\"all_results\"][k]\n",
    "                torchbeast_single_task_results[task][k].append(get_human_normalized_score(k, level_data))\n",
    "\n",
    "single_task_score_numpys = {k: task_dict_to_score_numpys(v) for k, v in torchbeast_single_task_results.items()}\n",
    "# logger.info(\"\\nCurriculum ablation aggregated scores\")\n",
    "# for task, scores in single_task_score_numpys.items():\n",
    "#     logger.info(f\"\\nSingle task: {task}\")\n",
    "#     print_aggregate_results(scores)\n",
    "# torchbeast_aggregated_single_task_results = defaultdict(list)\n",
    "# for key in valid_levels:\n",
    "#     task = folder_name_lookup[key].split(\"__\")[0]\n",
    "#     torchbeast_aggregated_single_task_results[key] = torchbeast_single_task_results[task][key][:3]\n",
    "# aggregated_single_task_score_numpys = task_dict_to_score_numpys(torchbeast_aggregated_single_task_results)\n",
    "\n",
    "aggregated_single_task_score_numpy = {k: v[k] for k, v in single_task_score_numpys.items()}\n",
    "logger.info(f\"\\nSingle task aggregated:\")\n",
    "print_aggregate_results(aggregated_single_task_score_numpy)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "\n",
    "tasks = [x.value.lower() for x in TASKS_BY_TASK_GROUP[AvalonTaskGroup.ALL]]\n",
    "single_task_matrix = np.zeros((len(tasks), len(tasks)))\n",
    "single_task_list_data = []\n",
    "for i, train_task in enumerate(tasks):\n",
    "    for j, test_task in enumerate(tasks):\n",
    "        mean_performance = metrics.aggregate_mean(single_task_score_numpys[train_task][test_task])\n",
    "        single_task_matrix[i][j] = mean_performance\n",
    "        single_task_list_data.append((train_task, test_task, mean_performance))\n",
    "\n",
    "df = pd.DataFrame(single_task_list_data, columns=[\"train_task\", \"test_task\", \"mean_performance\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "\n",
    "pivot_data = df.pivot(\"train_task\", \"test_task\", \"mean_performance\")\n",
    "fig, ax = plt.subplots(figsize=(10, 8))\n",
    "ax = sns.heatmap(pivot_data, annot=True, fmt=\".2f\", ax=ax, cbar=False)\n",
    "ax.set_title(\"Mean performance on each testing task for single task training\")\n",
    "ax.set_xlabel(\"Test task\")\n",
    "ax.set_ylabel(\"Train task\")\n",
    "\n",
    "plt.savefig(\"single_task_training_matrix.eps\", format=\"eps\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "\n",
    "## TORCHBEAST 500m learning curve\n",
    "learning_curve_checkpoints = {\n",
    "    \"bd10b76d-ce5b-4e0a-8514-47d75e777ae6\": [\n",
    "        \"model_step_3878400.tar\",\n",
    "        \"model_step_50310400.tar\",\n",
    "        \"model_step_100633600.tar\",\n",
    "        \"model_step_150387200.tar\",\n",
    "        \"model_step_199104000.tar\",\n",
    "        \"model_step_250752000.tar\",\n",
    "        \"model_step_301296000.tar\",\n",
    "        \"model_step_348435200.tar\",\n",
    "        \"model_step_399235200.tar\",\n",
    "        \"model_step_451862400.tar\",\n",
    "        \"final\",\n",
    "    ],\n",
    "    \"cdb06a43-4dfb-4d52-878d-8495075a4a38\": [\n",
    "        \"model_step_5440000.tar\",\n",
    "        \"model_step_48128000.tar\",\n",
    "        \"model_step_101523200.tar\",\n",
    "        \"model_step_149372800.tar\",\n",
    "        \"model_step_202572800.tar\",\n",
    "        \"model_step_250252800.tar\",\n",
    "        \"model_step_301126400.tar\",\n",
    "        \"model_step_349993600.tar\",\n",
    "        \"model_step_398956800.tar\",\n",
    "        \"model_step_450038400.tar\",\n",
    "        \"final\",\n",
    "    ],\n",
    "    \"42c4f566-49d1-45c6-a4f2-1185ac49238c\": [\n",
    "        \"model_step_4803200.tar\",\n",
    "        \"model_step_49753600.tar\",\n",
    "        \"model_step_101196800.tar\",\n",
    "        \"model_step_152124800.tar\",\n",
    "        \"model_step_198067200.tar\",\n",
    "        \"model_step_251043200.tar\",\n",
    "        \"model_step_299270400.tar\",\n",
    "        \"model_step_352342400.tar\",\n",
    "        \"model_step_399164800.tar\",\n",
    "        \"model_step_450905600.tar\",\n",
    "        \"final\",\n",
    "    ],\n",
    "}\n",
    "\n",
    "targets = [50_000_000 * i for i in range(11)]\n",
    "learning_curve_results = {i: defaultdict(list) for i in targets}\n",
    "learning_curve_actual_frames = defaultdict(list)\n",
    "for uuid, checkpoint_filenames in learning_curve_checkpoints.items():\n",
    "    for target, checkpoint_filename in zip(targets, checkpoint_filenames):\n",
    "        result_key = get_wandb_result_key(uuid, FIXED_WORLD_KEY, checkpoint_filename)\n",
    "        if checkpoint_filename == \"final\":\n",
    "            learning_curve_actual_frames[target].append(target)\n",
    "        else:\n",
    "            learning_curve_actual_frames[target].append(int(checkpoint_filename.split(\"_\")[-1].split(\".\")[0]))\n",
    "        try:\n",
    "            reported_data = json.loads(s3_client.load(result_key))\n",
    "        except KeyError:\n",
    "            logger.warning(f\"Missing result key {result_key}\")\n",
    "            continue\n",
    "        missing_levels = valid_levels - set(reported_data[\"all_results\"].keys())\n",
    "        if len(missing_levels) > 0:\n",
    "            logger.warning(f\"{result_key} is missing {len(missing_levels)} levels, ignoring incomplete data\")\n",
    "        else:\n",
    "            for k in valid_levels:\n",
    "                level_data = reported_data[\"all_results\"][k]\n",
    "                learning_curve_results[target][k].append(get_human_normalized_score(k, level_data))\n",
    "\n",
    "learning_curve_score_numpys = {k: task_dict_to_score_numpys(v) for k, v in learning_curve_results.items()}\n",
    "logger.info(\"\\nLearning curve scores\")\n",
    "for target, scores in learning_curve_score_numpys.items():\n",
    "    logger.info(f\"\\nSteps: {target}\")\n",
    "    print_aggregate_results(scores)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "\n",
    "\n",
    "def make_rows(latex_results: Dict, result_key: str, header: List):\n",
    "    rows = [header]\n",
    "    num_rows = len(PRINT_TASK_ORDER)\n",
    "    for i in range(num_rows):\n",
    "        task = PRINT_TASK_ORDER[i]\n",
    "        task_text = task.lower()\n",
    "        if task_text == \"simple\":\n",
    "            task_text = \"all basic\"\n",
    "        elif task_text == \"compositional\":\n",
    "            task_text = \"all comp.\"\n",
    "        row = [\"\\\\texttt{\" + task_text + \"}\"]\n",
    "        for key in latex_results:\n",
    "            row.append(latex_results[key][task][result_key])\n",
    "        rows.append(row)\n",
    "    return rows\n",
    "\n",
    "\n",
    "def drop_rows_by_columns(rows: List, ignore: List[int]):\n",
    "    rows_without_ignored_columns = []\n",
    "    for row in rows:\n",
    "        new_row = []\n",
    "        for i in range(len(row)):\n",
    "            if i in ignore:\n",
    "                continue\n",
    "            new_row.append(row[i])\n",
    "        rows_without_ignored_columns.append(new_row)\n",
    "    return rows_without_ignored_columns\n",
    "\n",
    "\n",
    "def print_latex_table(rows):\n",
    "    table = tabulate(rows, headers=\"firstrow\", tablefmt=\"latex_raw\")\n",
    "    table = table.replace(\"\\\\hline\", \"\\\\midrule\")\n",
    "    # manual fixes to get the right formatting\n",
    "    fixes = {\n",
    "        1: {\"text\": \"\\\\toprule\", \"replace\": True},\n",
    "        20: {\"text\": \"\\\\midrule\", \"replace\": False},\n",
    "        24: {\"text\": \"\\\\midrule\", \"replace\": False},\n",
    "        27: {\"text\": \"\\\\bottomrule\", \"replace\": True},\n",
    "    }\n",
    "    for i, line in enumerate(table.split(\"\\n\")):\n",
    "        if fixes.get(i):\n",
    "            logger.info(fixes[i][\"text\"])\n",
    "            if fixes[i].get(\"replace\"):\n",
    "                continue\n",
    "        logger.info(line)\n",
    "\n",
    "\n",
    "header = [\n",
    "    \"\\\\textbf{Task}\",\n",
    "    \"\\\\textbf{PPO}\",\n",
    "    \"\\\\textbf{Dreamer}\",\n",
    "    \"\\\\textbf{IMPALA}\",\n",
    "    \"\\\\textbf{IMPALA (500m)}\",\n",
    "    \"\\\\textbf{IMPALA (no curriculum)}\",\n",
    "]\n",
    "\n",
    "latex_results = {\n",
    "    \"ppo\": latex_aggregate_results(ppo_score_numpys),\n",
    "    \"dreamer\": latex_aggregate_results(dreamer_score_numpys),\n",
    "    \"impala\": latex_aggregate_results(torchbeast_score_numpys),\n",
    "    \"impala_500m\": latex_aggregate_results(curriciulum_score_numpys[\"500m\"]),\n",
    "    \"impala_no_curr\": latex_aggregate_results(curriciulum_score_numpys[\"no_curr\"]),\n",
    "}\n",
    "\n",
    "rows = make_rows(latex_results, \"mean\", header)\n",
    "print_latex_table(rows)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "\n",
    "\n",
    "rows = make_rows(latex_results, \"optimality_gap\", header)\n",
    "print_latex_table(rows)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "\n",
    "\n",
    "header = [\n",
    "    \"\\\\textbf{Task}\",\n",
    "    \"\\\\textbf{MT-TB}\",\n",
    "    \"\\\\textbf{MT-TA}\",\n",
    "    \"\\\\textbf{MT-TC}\",\n",
    "    \"\\\\textbf{ST-B}\",\n",
    "]\n",
    "\n",
    "latex_results = {\n",
    "    \"MT-TB\": latex_aggregate_results(torchbeast_score_numpys),\n",
    "    \"MT-TA\": latex_aggregate_results(curriciulum_score_numpys[\"multi_task_all\"]),\n",
    "    \"MT-TC\": latex_aggregate_results(curriciulum_score_numpys[\"multi_task_compositional\"]),\n",
    "    \"ST-B\": latex_aggregate_results(aggregated_single_task_score_numpy),\n",
    "}\n",
    "\n",
    "rows = make_rows(latex_results, \"mean\", header)\n",
    "print_latex_table(rows)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "\n",
    "\n",
    "rows = make_rows(latex_results, \"optimality_gap\", header)\n",
    "print_latex_table(rows)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "\n",
    "# CDF score plots\n",
    "\n",
    "rcParams[\"font.family\"] = \"serif\"\n",
    "rcParams[\"mathtext.default\"] = \"regular\"\n",
    "plt.rc(\"text\", usetex=False)\n",
    "\n",
    "\n",
    "def get_cdf(xs, a=0, b=10, n_bins=1000):\n",
    "    counts, bin_edges = np.histogram(xs, bins=np.linspace(a, b, n_bins + 1))\n",
    "    counts = counts / counts.sum()\n",
    "    cdf = counts[::-1].cumsum()[::-1]\n",
    "    # array of evenly-spaced scores for the x coords of the plot\n",
    "    bin_edges = bin_edges[:-1]\n",
    "    return bin_edges, cdf\n",
    "\n",
    "\n",
    "def get_tasks_in_sorted_order(data: Dict[str, np.ndarray]) -> List[str]:\n",
    "    tasks = list(data.keys())\n",
    "    tasks.sort(key=lambda x: get_cdf(data[x].flatten())[1][50], reverse=True)\n",
    "    return [t for t in tasks if t not in {\"simple\", \"compositional\", \"all\"}]\n",
    "\n",
    "\n",
    "def make_taskwise_cdf_plot(ax, data, title=\"\", no_ylabel=False, task_order=None):\n",
    "    if task_order is None:\n",
    "        tasks = get_tasks_in_sorted_order(data)\n",
    "    else:\n",
    "        tasks = task_order\n",
    "\n",
    "    bins, _ = get_cdf(data[\"eat\"].flatten())\n",
    "    cdfs = {}\n",
    "    for task in tasks:\n",
    "        _, cdfs[task] = get_cdf(data[task].flatten())\n",
    "\n",
    "    colors = [\n",
    "        hsv_to_rgb(c)\n",
    "        for c in [\n",
    "            (0.95, 0.7, 0.9),\n",
    "            (0.1, 0.8, 0.9),\n",
    "            (0.45, 0.8, 0.7),\n",
    "            (0.55, 0.7, 0.9),\n",
    "            (0.65, 0.75, 0.9),\n",
    "            (0.75, 0.5, 0.9),\n",
    "            (0, 0.1, 0.5),\n",
    "        ]\n",
    "    ]\n",
    "\n",
    "    for i, task in enumerate(tasks):\n",
    "        ax.plot(\n",
    "            bins,\n",
    "            cdfs[task],\n",
    "            color=colors[i // 3],\n",
    "            linestyle=[\"solid\", \"dashed\", \"dotted\"][i % 3],\n",
    "            lw=3,\n",
    "            label=task,\n",
    "            zorder=-i,\n",
    "        )  # f'$\\\\texttt{{{task}}}$')\n",
    "\n",
    "    ax.legend(ncol=2, loc=(0.53, 0.425), prop={\"family\": \"monospace\"})\n",
    "    # plt.setp(ax.legend().texts, family='typewriter')\n",
    "\n",
    "    ax.set_xlim((-0.05, 2))\n",
    "    ax.set_ylim((-0.03, 1.03))\n",
    "    ax.grid(linestyle=\"dotted\", linewidth=\"0.5\", color=(0.5, 0.5, 0.5))\n",
    "\n",
    "    fs = 17\n",
    "    ax.tick_params(labelsize=fs)\n",
    "    ax.set_xlabel(\"Human-normalized score $\\\\tau$\", fontsize=fs)\n",
    "    ax.set_ylabel(\"Fraction of runs with score $\\\\geq \\\\tau$\", fontsize=fs) if not no_ylabel else None\n",
    "    ax.set_yticklabels([]) if no_ylabel else None\n",
    "    ax.set_title(title, fontsize=fs, pad=15)\n",
    "    ax.spines[\"right\"].set_visible(False)\n",
    "    ax.spines[\"top\"].set_visible(False)\n",
    "\n",
    "\n",
    "def make_aggregate_cdf_plot(ax, datasets, names, title=\"\", no_ylabel=False):\n",
    "    # colors = get_colors(len(names))\n",
    "    # colors = [.9,.2,0], [0,.5,.8]\n",
    "    colors = hsv_to_rgb([0.95, 0.7, 0.9]), hsv_to_rgb([0.65, 0.75, 0.9]), hsv_to_rgb((0.45, 0.8, 0.7))\n",
    "\n",
    "    bins, _ = get_cdf(datasets[0][\"eat\"].flatten())\n",
    "\n",
    "    # for x in [0,.5,1,1.5,2]:\n",
    "    #   ax.axvline(x=x, color=(.5,.5,.5), linestyle='dotted', lw=1, zorder=-1000)\n",
    "    # for y in [0,.2,.4,.6,.8,1]:\n",
    "    #   ax.axhline(y=y, color=(.5,.5,.5), linestyle='dotted', lw=1, zorder=-1000)\n",
    "\n",
    "    # cdfs = {}\n",
    "    for i, (data, name) in enumerate(zip(datasets, names)):\n",
    "        for j, task in enumerate([\"simple\", \"compositional\", \"all\"]):\n",
    "            _, cdf = get_cdf(datasets[i][task].flatten())\n",
    "            ax.plot(\n",
    "                bins,\n",
    "                cdf,\n",
    "                color=colors[i],\n",
    "                label=f'{name} ({task if task != \"simple\" else \"basic\"})',\n",
    "                lw=3,\n",
    "                linestyle=[\"dotted\", \"dashed\", \"solid\"][j],\n",
    "                zorder=-i,\n",
    "            )\n",
    "\n",
    "    ax.legend(ncol=1, prop={\"family\": \"monospace\"})\n",
    "\n",
    "    ax.set_xlim((-0.05, 2))\n",
    "    ax.set_ylim((-0.03, 1.03))\n",
    "    ax.grid(linestyle=\"dotted\", linewidth=\"0.5\", color=(0.5, 0.5, 0.5))\n",
    "\n",
    "    fs = 17\n",
    "    ax.tick_params(labelsize=fs)\n",
    "    ax.set_xlabel(\"Human-normalized score $\\\\tau$\", fontsize=fs)\n",
    "    ax.set_ylabel(\"Fraction of runs with score $\\geq \\\\tau$\", fontsize=fs) if not no_ylabel else None\n",
    "    ax.set_yticklabels([]) if no_ylabel else None\n",
    "    ax.set_title(title, fontsize=fs, pad=15)\n",
    "    ax.spines[\"right\"].set_visible(False)\n",
    "    ax.spines[\"top\"].set_visible(False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "\n",
    "rcParams[\"font.family\"] = \"serif\"\n",
    "rcParams[\"mathtext.default\"] = \"regular\"\n",
    "plt.rc(\"text\", usetex=False)\n",
    "\n",
    "fig, axs = plt.subplots(ncols=2, nrows=2, figsize=(12, 10))\n",
    "\n",
    "task_order = get_tasks_in_sorted_order(torchbeast_score_numpys)\n",
    "make_taskwise_cdf_plot(axs[0][0], torchbeast_score_numpys, title=\"IMPALA\", task_order=task_order)\n",
    "make_taskwise_cdf_plot(axs[0][1], ppo_score_numpys, title=\"PPO\", no_ylabel=True, task_order=task_order)\n",
    "make_taskwise_cdf_plot(axs[1][0], dreamer_score_numpys, title=\"Dreamer\", no_ylabel=True, task_order=task_order)\n",
    "make_aggregate_cdf_plot(\n",
    "    axs[1][1],\n",
    "    [torchbeast_score_numpys, ppo_score_numpys, dreamer_score_numpys],\n",
    "    names=[\"IMPALA\", \"PPO\", \"Dreamer\"],\n",
    "    title=\"Aggregate\",\n",
    "    no_ylabel=True,\n",
    ")\n",
    "\n",
    "fig.tight_layout()\n",
    "\n",
    "plt.savefig(\"taskwise_agent_scores.eps\", format=\"eps\")\n",
    "# plt.savefig(\"taskwise_agent_scores.pdf\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "\n",
    "ts = sorted(learning_curve_score_numpys.keys())\n",
    "true_ts = [np.mean(learning_curve_actual_frames[k]) for k in ts]\n",
    "tasks = [t for t in learning_curve_score_numpys[ts[0]].keys() if t not in [\"simple\", \"compositional\", \"all\"]]\n",
    "curves = {task: [learning_curve_score_numpys[t][task].mean() for t in ts] for task in tasks}\n",
    "\n",
    "# order tasks\n",
    "# fig, ax = plt.subplots(figsize=(8, 6))\n",
    "\n",
    "\n",
    "def plot_all_tasks(ax):\n",
    "    colors = [\n",
    "        hsv_to_rgb(c)\n",
    "        for c in [\n",
    "            (0.95, 0.7, 0.9),\n",
    "            (0.1, 0.8, 0.9),\n",
    "            (0.45, 0.8, 0.7),\n",
    "            (0.55, 0.7, 0.9),\n",
    "            (0.65, 0.75, 0.9),\n",
    "            (0.75, 0.5, 0.9),\n",
    "            (0, 0.1, 0.5),\n",
    "        ]\n",
    "    ]\n",
    "    for i, task in enumerate(task_order):\n",
    "        ax.plot(\n",
    "            true_ts,\n",
    "            curves[task],\n",
    "            color=colors[i // 3],\n",
    "            linestyle=[\"solid\", \"dashed\", \"dotted\"][i % 3],\n",
    "            lw=3,\n",
    "            label=task,\n",
    "            zorder=-i,\n",
    "        )\n",
    "\n",
    "    ax.legend(ncol=3, prop={\"family\": \"monospace\"})\n",
    "\n",
    "    # ax.set_xlim((-.05,2))\n",
    "    ax.set_ylim((-0.03, 1.03))\n",
    "    ax.grid(linestyle=\"dotted\", linewidth=\"0.5\", color=(0.5, 0.5, 0.5))\n",
    "    fs = 17\n",
    "    ax.tick_params(labelsize=fs)\n",
    "    ax.set_xlabel(\"Environment steps\", fontsize=fs)\n",
    "    ax.set_ylabel(\"Mean performance on task\", fontsize=fs)\n",
    "    ax.spines[\"right\"].set_visible(False)\n",
    "    ax.spines[\"top\"].set_visible(False)\n",
    "\n",
    "\n",
    "fig, ax = plt.subplots(figsize=(8, 5))\n",
    "plot_all_tasks(ax)\n",
    "plt.savefig(\"taskwise_learning_curves.eps\", format=\"eps\")"
   ]
  }
 ],
 "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.9.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}