{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "WfNbd9F5aKN5"
      },
      "source": [
        "# Imports and Styling"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Yg2sXPcUfrld",
        "outputId": "31eee514-70d8-4d63-f6a1-d1c246508041"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Looking in indexes: https://pypi.org/simple, https://us-python.pkg.dev/colab-wheels/public/simple/\n",
            "Collecting colabtools\n",
            "  Downloading colabtools-0.0.1-py3-none-any.whl (14 kB)\n",
            "Installing collected packages: colabtools\n",
            "Successfully installed colabtools-0.0.1\n"
          ]
        }
      ],
      "source": [
        "#@title Imports and styling\n",
        "from typing import List, Tuple\n",
        "import pdb\n",
        "!pip install colabtools\n",
        "from collections import defaultdict\n",
        "import enum\n",
        "import functools\n",
        "import os\n",
        "import random\n",
        "import sys\n",
        "import time\n",
        "from datetime import datetime\n",
        "from absl import flags\n",
        "from copy import deepcopy\n",
        "import matplotlib\n",
        "import matplotlib.pyplot as plt\n",
        "import numpy as np\n",
        "import numpy.ma as ma\n",
        "import pandas as pd\n",
        "import plotnine as gg\n",
        "import seaborn as sns\n",
        "from numpy.linalg import inv"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "KvtCk8yqRfvM",
        "outputId": "caf1ca92-bb82-4cd1-8413-9b717d572cbb"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Looking in indexes: https://pypi.org/simple, https://us-python.pkg.dev/colab-wheels/public/simple/\n",
            "Collecting envlogger\n",
            "  Downloading envlogger-1.0.8-cp38-cp38-manylinux2014_x86_64.whl (5.5 MB)\n",
            "\u001b[K     |████████████████████████████████| 5.5 MB 4.8 MB/s \n",
            "\u001b[?25hRequirement already satisfied: numpy in /usr/local/lib/python3.8/dist-packages (from envlogger) (1.21.6)\n",
            "Requirement already satisfied: setuptools!=50.0.0 in /usr/local/lib/python3.8/dist-packages (from envlogger) (57.4.0)\n",
            "Requirement already satisfied: absl-py in /usr/local/lib/python3.8/dist-packages (from envlogger) (1.3.0)\n",
            "Requirement already satisfied: protobuf>=3.14 in /usr/local/lib/python3.8/dist-packages (from envlogger) (3.19.6)\n",
            "Collecting dm-env\n",
            "  Downloading dm_env-1.6-py3-none-any.whl (26 kB)\n",
            "Requirement already satisfied: dm-tree in /usr/local/lib/python3.8/dist-packages (from dm-env->envlogger) (0.1.7)\n",
            "Installing collected packages: dm-env, envlogger\n",
            "Successfully installed dm-env-1.6 envlogger-1.0.8\n"
          ]
        }
      ],
      "source": [
        "!pip install envlogger"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ILy1KYTXT7l4"
      },
      "outputs": [],
      "source": [
        "from envlogger.converters import codec\n",
        "from envlogger.proto import storage_pb2"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "7m6nnyd1DJDf"
      },
      "outputs": [],
      "source": [
        "import seaborn as sn\n",
        "sn.set(color_codes=True)\n",
        "sn.set(style='ticks')\n",
        "plt.style.use('classic')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "r05HA2UWhKz6"
      },
      "source": [
        "# Utilities\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "m9dmhPOZ3qnt"
      },
      "outputs": [],
      "source": [
        "#@title Misc.\n",
        "def indices(a, func):\n",
        "    return [i for (i, val) in enumerate(a) if func(val)]\n",
        "\n",
        "def sample_mult(p):\n",
        "  \"\"\" sample from a multinomial distribution with probabilities p\"\"\"\n",
        "  multinom = np.random.multinomial(1, p)\n",
        "  indc1 = indices(multinom, lambda x: x>0)\n",
        "  return indc1[0]\n",
        "\n",
        "def noise():\n",
        "  noise_ = 0.1 * np.random.normal()\n",
        "  return noise_\n",
        "\n",
        "def _unit_test_P(P):\n",
        "  assert np.allclose(P.sum(axis=1), 1), 'P matrix is not stochastic!'  "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "UdP8yGbwIcyJ"
      },
      "outputs": [],
      "source": [
        "#@title Random MDP Generation \n",
        "def _unit_test_P(P):\n",
        "  assert np.allclose(P.sum(axis=1), 1), 'P matrix is not stochastic!'\n",
        "\n",
        "def randomMDP(num_states=10, k=5):\n",
        "  \"\"\"\n",
        "  For each state-action pair (s,a), the distribution over the next state, P(s,a,·), is determined\n",
        "  by choosing 5 non-zero entries uniformly from all  10  states, filling  these  5  entries  with\n",
        "  values  uniformly drawn  from [0,1],  and  finally  normalizing P(s,a,·).  \n",
        "\n",
        "  The mean rewards were likewise sampled uniformly and independently from [0,1],\n",
        "  and the actual reward signals have additive Gaussian noise with standard deviation 0.1.  \n",
        "\n",
        "  \"\"\"\n",
        "  if k<0:\n",
        "    P = np.random.uniform(0, 1, (num_states, num_states))\n",
        "    R = np.random.uniform(0, 1, (num_states, 1))\n",
        "    z = np.sum(P,1, keepdims=True)\n",
        "    P = P / z\n",
        "  else:\n",
        "    P0 = np.random.uniform(0, 1, (num_states, k))\n",
        "    P = np.zeros((num_states, num_states))\n",
        "    for i in range(0, num_states):\n",
        "      rp = np.random.permutation(num_states)\n",
        "      idx = rp[0:k]\n",
        "      P[i, idx] = P0[i, :]\n",
        "    R = np.random.uniform(0, 1, (num_states, 1))\n",
        "    z = np.sum(P,1, keepdims=True)\n",
        "    P = P / z\n",
        "  return   P, R"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "bLKE2vyOtIVc"
      },
      "outputs": [],
      "source": [
        " #@title  sample_trajectories_transitions\n",
        "\n",
        "def sample_trajectories_transitions(\n",
        "    P1, R1, P2, R2, num_trajectories=5, max_trajectory_length=5,\n",
        "    num_actions=2, startstate=None, mu0 = None, seed=None):\n",
        "  \"\"\"Takes transitions samples from an environment.\n",
        "\n",
        "  Args:\n",
        "    P: mdp transition model\n",
        "    R: mdp reward model\n",
        "    num_trajectories: The total number of trajectories to sample.\n",
        "    max_trajectory_length: The maximum length of the trajectory.\n",
        "    seed: seed for randomness \n",
        "  Returns:\n",
        "    The trajectories collected from the environment:\n",
        "      This is a 4-tuple containing the batch of \n",
        "      state, action, reward, next-state.\n",
        "  \"\"\"\n",
        "\n",
        "  if mu0 is None:\n",
        "    mu0 = np.ones((R1.shape[0])) / R1.shape[0] \n",
        "\n",
        "  trajectory_data = np.zeros((num_trajectories * max_trajectory_length, 4))\n",
        "  trajectory_length = 0\n",
        "\n",
        "\n",
        "  for i in range(num_trajectories):\n",
        "    if startstate is None:\n",
        "      s_o = sample_mult(mu0)\n",
        "    else:\n",
        "      s_o = startstate # fix start state\n",
        "    trajectory_data[(i) * max_trajectory_length, 0] = s_o\n",
        "    for j in range(max_trajectory_length):\n",
        "      idx = (i) * max_trajectory_length + j\n",
        "      s_t = int(trajectory_data[idx, 0])\n",
        "      a_t = np.random.randint(num_actions)\n",
        "      trajectory_data[idx, 1] = a_t \n",
        "\n",
        "      if a_t == 0:\n",
        "        r = R1[s_t, 0]\n",
        "        p = P1[s_t, :]\n",
        "      else:\n",
        "        r = R2[s_t, 0]\n",
        "        p = P2[s_t, :]\n",
        "      \n",
        "      trajectory_data[idx, 2] = r + noise()\n",
        "      trajectory_data[idx, 3] = sample_mult(p)\n",
        "      \n",
        "\n",
        "      if j < max_trajectory_length-1:\n",
        "          trajectory_data[idx+1, 0] = trajectory_data[idx, 3]\n",
        "\n",
        "  return trajectory_data"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5CH8dyIfcmDf"
      },
      "source": [
        "learn_model_from_trajectories: To learn the model from the experience-data, we consider the count based approach:\n",
        "\n",
        "Transition Dynamics computation $\\hat{P^t}(s'|s,a) = \\frac{\\texttt{count}(s,a,s')}{m} = \\frac{\\sum_{i=1}^k \\mathrm{1}\\{ s^{'}_i=s'\\}} {k}, \\texttt{where } \\{ s'_i, t_i\\} \\sim P^t_{\\texttt{true}}(\\cdot|s,a) \\forall 1 \\leq i \\leq m.$\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "7-3WFiT2369B",
        "cellView": "form"
      },
      "outputs": [],
      "source": [
        "#@title learn_model_from_trajectories\n",
        " \n",
        "def learn_model_from_trajectories(\n",
        "    num_states, \n",
        "    num_actions, \n",
        "    trajectories, \n",
        "    averageinit=True,\n",
        "    alpha=None,\n",
        "    P1Prior=None,\n",
        "    P2Prior=None):\n",
        "  \"\"\"Learns a transition dynamics from sampled data.\n",
        "\n",
        "  Args:\n",
        "    mdp: The MDP to sample from.\n",
        "    mdp_wall_locs: The wall locations in that mdp.\n",
        "    num_trajectories: The total number of trajectories to sample.\n",
        "    max_trajectory_length: The maximum length of the trajectory.\n",
        "    averageinit: boolean true if initialization choice is uniform.\n",
        "  Returns:\n",
        "    The transition model P1, P2\n",
        "    The reward model R1, R2 per action\n",
        "  \"\"\"\n",
        "\n",
        "  if (averageinit):\n",
        "    P1 = np.ones((num_states, num_states))/num_states\n",
        "    P2 = np.ones((num_states, num_states))/num_states\n",
        "    R1 = 0.5 * np.ones((num_states, 1))\n",
        "    R2 = 0.5 * np.ones((num_states, 1)) \n",
        "  else:\n",
        "    P1 = np.zeros((num_states, num_states))\n",
        "    P2 = np.zeros((num_states, num_states))\n",
        "    R1 = np.zeros((num_states, 1))\n",
        "    R2 = np.zeros((num_states, 1)) \n",
        "\n",
        "\n",
        "  if type(alpha) is not list: alpha = [alpha, alpha]\n",
        "\n",
        "   \n",
        "  for s in range(0, num_states):\n",
        "    idx1 = (trajectories[:,0] == s) & (trajectories[:,1] == 0)\n",
        "    samples_count = np.sum(idx1)\n",
        "    if samples_count > 0:\n",
        "      P1[s, :], _ = np.histogram(trajectories[idx1,3].astype(int), bins=range(num_states+1))\n",
        "      P1[s, :] = P1[s, :]/samples_count\n",
        "      R1[s,0] = np.mean(trajectories[idx1, 2])\n",
        "\n",
        "\n",
        "    idx2 = (trajectories[:,0] == s) & (trajectories[:,1] == 1)\n",
        "    samples_count2 = np.sum(idx2)       \n",
        "    if samples_count2 > 0:\n",
        "      P2[s, :], _ = np.histogram(trajectories[idx2,3].astype(int), bins=range(num_states+1))\n",
        "      P2[s, :] = P2[s, :]/samples_count2\n",
        "      R2[s,0] = np.mean(trajectories[idx2, 2])\n",
        "\n",
        "  if np.any(alpha) is not None:\n",
        "    _unit_test_P(P1Prior)\n",
        "    _unit_test_P(P2Prior)\n",
        "    P_1 = (1-alpha[0]) * P1 + (alpha[0])*P1Prior\n",
        "    P_2 = (1-alpha[1]) * P2 + (alpha[1])*P2Prior\n",
        "  else:\n",
        "    P_1 = P1\n",
        "    P_2 = P2\n",
        "\n",
        "  _unit_test_P(P_1)\n",
        "  _unit_test_P(P_2)  \n",
        " \n",
        "  return P_1, R1, P_2, R2"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "fQD6ZTqU3830",
        "cellView": "form"
      },
      "outputs": [],
      "source": [
        " #@title dynamic programming\n",
        "def polPR(P1, R1, P2, R2, pi):\n",
        "  P = np.zeros((P1.shape[0], P1.shape[1]))\n",
        "  R = np.zeros((R1.shape[0], R1.shape[1]))\n",
        "  for i in range(0, len(pi)):\n",
        "    P[i, :] = (1 - pi[i,0]) * P1[i, :] + pi[i,0] * P2[i, :]\n",
        "    R[i, 0] = (1 - pi[i,0]) * R1[i, 0] + pi[i,0] * R2[i,0]\n",
        "  return P, R\n",
        "\n",
        "\n",
        "\n",
        "def policyeval(pi, P1, R1, P2, R2, gamma):\n",
        "  p, r = polPR(P1, R1, P2, R2, pi)\n",
        "  v = np.dot(inv((np.eye(p.shape[0]) - gamma*p)), r)\n",
        "  return v\n",
        "\n",
        "\n",
        "\n",
        "def dynamicprogramming(P1, R1, P2, R2, gamma):\n",
        "  N = R1.shape[0]\n",
        "  pol = np.ones((N, 1))\n",
        "  pol0 = np.zeros((N,1))\n",
        "  t = 0\n",
        "  max_iterations = 100\n",
        "\n",
        "  while (np.linalg.norm(pol-pol0, 2) > 0):\n",
        "    pol0 = deepcopy(pol)\n",
        "    V = policyeval(pol, P1, R1, P2, R2, gamma)\n",
        "    for i in range(0, N):\n",
        "      q1 = R1[i,0] + gamma * np.dot(P1[i,:], V[:,0])\n",
        "      q2 = R2[i,0] + gamma * np.dot(P2[i,:], V[:,0])\n",
        "      if (q1 > q2):\n",
        "        pol[i,0] = 0\n",
        "      else:\n",
        "        pol[i,0] = 1\n",
        "    t = t+1\n",
        "    if (t > max_iterations):\n",
        "      pol, V = vi(P1, R1, P2, R2, gamma, V)\n",
        "      return pol, V\n",
        "  return pol, V\n",
        "\n",
        "\n",
        "\n",
        "def vi(P1, R1, P2, R2, gamma, V=None):\n",
        "  if V is not None:\n",
        "    V = R1/ (1-gamma)\n",
        "\n",
        "  N = R1.shape[0]\n",
        "  V0 = np.zeros((N, 1))\n",
        "  max_iterations = 100\n",
        "\n",
        "  while (np.linalg.norm(V-V0, 2) > 1e-5):\n",
        "    V0 = deepcopy(V)\n",
        "    for s in range(0, N):\n",
        "      q1 = R1[s,0] + gamma * np.dot(P1[s,:], V0[:,0])\n",
        "      q2 = R2[s,0] + gamma * np.dot(P2[s,:], V0[:,0])\n",
        "      V[s,0] = max(q1, q2)   \n",
        "\n",
        "  policy =  np.zeros((N, 1))\n",
        "  for s in range(0, N):\n",
        "    q1 = R1[s,0] + gamma * np.dot(P1[s,:], V[:,0])\n",
        "    q2 = R2[s,0] + gamma * np.dot(P2[s,:], V[:,0])\n",
        "    if (q1 > q2):\n",
        "      policy[s, 0] = 0\n",
        "    else:\n",
        "      policy[s, 0] = 1\n",
        "  \n",
        "  return policy, V"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "fx9JS6g938kn"
      },
      "outputs": [],
      "source": [
        " #@title processTrajectory\n",
        "\n",
        " def processTrajectory(P1, R1, P2, R2, samples, gammas, gammaEval,\n",
        "                       trials=None, alpha=None, P1o=None, P2o=None):\n",
        "  if trials is None:\n",
        "    trials = 400\n",
        "  if np.isinf(samples):\n",
        "    trials = 2\n",
        "  \n",
        "  max_trajectory_length = 10\n",
        "  N = R1.shape[0]\n",
        "\n",
        "  res = np.zeros((trials, len(gammas)))\n",
        "  for j in range(0, trials):\n",
        "    if (not np.isinf(samples)):\n",
        "      trs_trajectories =  sample_trajectories_transitions(P1, R1, P2, R2, samples, \n",
        "                                                          max_trajectory_length)\n",
        "    for k in range(0, len(gammas)):\n",
        "      if (not np.isinf(samples)):\n",
        "        Phat_1, Rhat_1, Phat_2, Rhat_2 = learn_model_from_trajectories(num_states=N, num_actions=2,\n",
        "                                                                       trajectories=trs_trajectories,\n",
        "                                                                       averageinit=True,\n",
        "                                                                       alpha=alpha,\n",
        "                                                                       P1Prior=P1o,\n",
        "                                                                       P2Prior=P2o)\n",
        "      else:\n",
        "        Phat_1 = P1\n",
        "        Rhat_1 = R1\n",
        "        Phat_2 = P2\n",
        "        Rhat_2 = R2\n",
        "\n",
        "      gamma =  gammas[k]\n",
        "      policy, _ = dynamicprogramming(Phat_1, R1, Phat_2, R2, gamma)\n",
        "      # Uncomment below if learning rewards too.\n",
        "      # policy, _ = dynamicprogramming(Phat_1, Rhat_1, Phat_2, Rhat_2, gamma)\n",
        "      res[j,k] = np.mean(policyeval(policy, P1, R1, P2, R2, gammaEval), axis=0)\n",
        "\n",
        "\n",
        "  y = np.mean(res, 0).reshape((1, len(gammas)))\n",
        "  err = np.sqrt(np.var(res, 0) / (trials-1) ).reshape((1, len(gammas)))\n",
        "  templist = [np.sum( np.abs(Phat_1 - P1), axis=1)]\n",
        "  templist.append(np.sum( np.abs(Phat_2 - P2), axis=1))\n",
        "  et = np.max(templist) \n",
        "  return y, err, et"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "AZLAMFTpJKxd"
      },
      "source": [
        "# Planning with online meta-learned prior \n",
        "\n",
        "$\\Big|\\Big|V^{\\pi^*_{M,\\gamma_{\\texttt{eval}}}}_{M,\\gamma_{\\texttt{eval}}} - V^{\\pi^{*}_{\\hat{M},\\gamma}}_{M, \\gamma_{\\texttt{eval}}} \\Big|\\Big|_{\\infty}$"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "UKQWla-y9IkU"
      },
      "outputs": [],
      "source": [
        "# @title plotloss, plotoptgamma\n",
        "from fractions import Fraction\n",
        "def plotloss(y, e, gammas, numtasks):\n",
        "  \n",
        "  fig, ax1 = plt.subplots()\n",
        "  lines = ['d', 'x', 'o', 's']\n",
        "  for i in range(0, numtasks):\n",
        "    if i<5 or i%3== 0:\n",
        "      yk = y[i,:].reshape((1,y.shape[-1]))\n",
        "      ek = e[i,:].reshape((1,y.shape[-1]))\n",
        "      plt.errorbar(gammas, yk[0,:], yerr=ek[0,:],\n",
        "                  label = 'T={}'.format(int(i+1)))\n",
        "  ax1.set_xlabel('$\\gamma$')\n",
        "  ax1.set_ylabel('Planning Loss')\n",
        "  ax1.set_xticks(gammas)\n",
        "  ax1.legend(loc='upper center', bbox_to_anchor=(0.5, -0.19),\n",
        "            fancybox=True, shadow=True, ncol=7,\n",
        "            facecolor='w', fontsize=10)\n",
        "  plt.title('Planning w/ Online Meta-Learning')\n",
        "  matplotlib.rc('axes', edgecolor='black')\n",
        "  plt.show()    \n",
        "  return fig\n",
        "\n",
        "def plotoptgamma(y, e, numtasks):\n",
        "  fig, ax1 = plt.subplots()\n",
        "  xtoplot = list(range(1, numtasks+1))\n",
        "  plt.errorbar(xtoplot,\n",
        "               y.reshape((1,y.shape[-1]))[0,:], \n",
        "               yerr=e.reshape((1,y.shape[-1]))[0,:])\n",
        "  ax1.set_xlabel('Number of task seen (t)')\n",
        "  ax1.set_ylabel('Optimal Gamma ($\\gamma^*$)')\n",
        "  ax1.set_xticks(xtoplot)\n",
        "  ax1.legend(loc='upper center', bbox_to_anchor=(0.5, -0.19),\n",
        "            fancybox=True, shadow=True, ncol=7,\n",
        "            facecolor='w', fontsize=10)\n",
        "  plt.title('Dependence of effective planning horizon on Tasks.')\n",
        "  matplotlib.rc('axes', edgecolor='black')\n",
        "  plt.show()    \n",
        "  return fig"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "3cDjeeeXf0YU",
        "cellView": "form"
      },
      "outputs": [],
      "source": [
        "# @title processwithintask\n",
        "def processwithintask(P1, R1, P2, R2, samples, gammas, gammaEval,\n",
        "                      trials=None, alpha=None, P1o=None, P2o=None):\n",
        "  '''\n",
        "  Computes planning loss for round t using meta-initialization \n",
        "  of prior mean and learning rate. \n",
        "\n",
        "  Steps:  \n",
        "  - generate trajectories from Pt with low sample size samples*max_trajectory_length\n",
        "  - estimate Pt_hat using current estimate of Prior Po_hat\n",
        "  - calculate planning loss for each gamma\n",
        "  \n",
        "  Args:\n",
        "      P1, R1, P2, R2: The MDP to sample from.\n",
        "      samples: dataset size per task.\n",
        "      gammas: list of guidance discount factors.\n",
        "      gammaEval: true gamma for evaluation.\n",
        "      trials: number of seeds to run this function. \n",
        "      alpha: current learning rate\n",
        "      P1o, P2o: current estimate of prior mean.\n",
        "  \n",
        "  Returns: planningloss and std-error for task t averaged over #trials\n",
        "  '''\n",
        "  if trials is None:\n",
        "    trials = 400\n",
        "  \n",
        "  max_trajectory_length = 10\n",
        "  N = R1.shape[0]\n",
        "\n",
        "  res = np.zeros((trials, len(gammas)))\n",
        "  for j in range(0, trials):\n",
        "    trs_trajectories =  sample_trajectories_transitions(P1, R1, P2, R2, samples,\n",
        "                                                        max_trajectory_length)\n",
        "    for k in range(0, len(gammas)):\n",
        "      Phat_1, Rhat_1, Phat_2, Rhat_2 = learn_model_from_trajectories(num_states=N, num_actions=2,\n",
        "                                                                     trajectories=trs_trajectories,\n",
        "                                                                     averageinit=True,\n",
        "                                                                     alpha=alpha,\n",
        "                                                                     P1Prior=P1o,\n",
        "                                                                     P2Prior=P2o)\n",
        "      _unit_test_P(Phat_1)\n",
        "      _unit_test_P(Phat_2)\n",
        "\n",
        "      gamma =  gammas[k]\n",
        "      policy, _ = dynamicprogramming(Phat_1, R1, Phat_2, R2, gamma)\n",
        "      res[j,k] = np.mean(policyeval(policy, P1, R1, P2, R2, gammaEval), axis=0)\n",
        "\n",
        "  y = np.mean(res, 0).reshape((1, len(gammas)))\n",
        "  err = np.sqrt(np.var(res, 0) / (trials-1) ).reshape((1, len(gammas)))\n",
        "  templist = [np.sum( np.abs(Phat_1 - P1), axis=1)]\n",
        "  templist.append(np.sum( np.abs(Phat_2 - P2), axis=1))\n",
        "  et = np.max(templist) \n",
        "\n",
        "  # calculating Vpistar_Mgammeval\n",
        "  policy_star, _ = dynamicprogramming(P1, R1, P2, R2, trueGamma)\n",
        "  VMAX = np.mean(policyeval(policy_star, P1, R1, P2, R2, trueGamma), axis=0) \n",
        "\n",
        "  # calculate loss for this particular Pt Po run. \n",
        "  loss = VMAX - y\n",
        "  err = err\n",
        "  \n",
        "  return loss, err  "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "xtRRA-y2wbvf"
      },
      "outputs": [],
      "source": [
        "# @title update_prior_online\n",
        "def update_prior_online(P1ohatprev, P2ohatprev, roundt, m,\n",
        "                        P1, P2, R1, R2, max_trajectory_length,\n",
        "                        trials):\n",
        "  '''\n",
        "  Args:\n",
        "      P1ohatprev: Estimate of prior mean Po1 at prev round t-1\n",
        "      P2ohatprev: Estimate of prior mean Po2 at prev round t-1.\n",
        "      t: round t.\n",
        "      m: number of samples from current task.\n",
        "  \n",
        "  Returns:\n",
        "      P1ohat: estimate of prior mean p1o at round t.\n",
        "      P2ohat: estimate of prior mean p2o at round t.\n",
        "  '''\n",
        "  num_states = P1.shape[0]\n",
        "\n",
        "  Pres1 = np.zeros((trials, num_states, num_states))\n",
        "  Pres2 = np.zeros((trials, num_states, num_states))\n",
        "  P1_avgm_roundt = np.zeros((num_states, num_states))\n",
        "  P2_avgm_roundt = np.zeros((num_states, num_states))\n",
        "\n",
        "  P1ohat = np.zeros((num_states, num_states))\n",
        "  P2ohat = np.zeros((num_states, num_states))\n",
        "\n",
        "  sa_counts = np.zeros((trials, num_states, 2))\n",
        "\n",
        "  for j in range(0, trials):\n",
        "    trajectories =  sample_trajectories_transitions(P1, R1, P2, R2, m,\n",
        "                                                        max_trajectory_length)\n",
        "    \n",
        "\n",
        "    for s in range(0, num_states):\n",
        "      idx1 = (trajectories[:,0] == s) & (trajectories[:,1] == 0)\n",
        "      samples_count = np.sum(idx1)\n",
        "      if samples_count > 0:\n",
        "        Pres1[j, s, :], _ = np.histogram(trajectories[idx1,3].astype(int), bins=range(num_states+1))\n",
        "        Pres1[j, s, :] = Pres1[j, s, :]/samples_count      \n",
        "      sa_counts[j,s,0] = samples_count\n",
        "\n",
        "      idx2 = (trajectories[:,0] == s) & (trajectories[:,1] == 1)\n",
        "      samples_count2 = np.sum(idx2)       \n",
        "      if samples_count2 > 0:\n",
        "        Pres2[j, s, :], _ = np.histogram(trajectories[idx2,3].astype(int), bins=range(num_states+1))\n",
        "        Pres2[j, s, :] = Pres2[j, s, :]/samples_count2   \n",
        "      sa_counts[j,s,1] = samples_count2\n",
        "  \n",
        "  P1_avgm_roundt = np.mean(Pres1, 0)\n",
        "  P2_avgm_roundt = np.mean(Pres2, 0)\n",
        "  \n",
        "  sa_counts_avg = np.mean(sa_counts, 0)\n",
        "\n",
        "  P1_avgm_roundt = P1_avgm_roundt/np.sum(P1_avgm_roundt, 1, keepdims=True)\n",
        "  P2_avgm_roundt = P2_avgm_roundt/np.sum(P2_avgm_roundt, 1, keepdims=True)\n",
        "  \n",
        "  _unit_test_P(P1_avgm_roundt)\n",
        "  _unit_test_P(P2_avgm_roundt)\n",
        "\n",
        "  # prior at round t <--- (1-1/t) * estimate at round (t-1) + (1/t) * counting avg of m samples in round t\n",
        "  P1ohat =  (1-1/roundt)* P1ohatprev + (1/roundt)*  P1_avgm_roundt\n",
        "  P2ohat =  (1-1/roundt)* P2ohatprev + (1/roundt)*  P2_avgm_roundt\n",
        "  \n",
        "  _unit_test_P(P1ohat)\n",
        "  _unit_test_P(P2ohat)\n",
        "  return P1ohat, P2ohat, P1_avgm_roundt, P2_avgm_roundt, sa_counts_avg\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "0j0BB6nkwgqx"
      },
      "outputs": [],
      "source": [
        "# @title update_variance_online\n",
        "def update_variance_online(sigmao1hat_t, sigmao2hat_t,\n",
        "                           P1ohat_t, P2ohat_t,\n",
        "                           P1ohat_tp1, P2ohat_tp1,\n",
        "                           P1_avgm_roundt, P2_avgm_roundt,\n",
        "                           roundt):\n",
        "\n",
        "  '''\n",
        "  Welford's online algorithm\n",
        "\n",
        "  Args:\n",
        "      sigmao1hat_t: estimate of prior variance sigma0^2 at round t (for action1).\n",
        "      sigmao2hat_t: estimate of prior variance sigma0^2 at round t (for action2).\n",
        "      P1ohat_t, P2ohat_t: mean of prior at round t\n",
        "      P1ohat_t, P2ohat_t: mean of prior at round t+1\n",
        "      P1_avgm_roundt, P2_avgm_roundt: average of m samples from current task\n",
        "  \n",
        "  Returns:\n",
        "      sigmao1hat_tp1: estimate of prior variance sigma0^2 at round t+1 (for action1).\n",
        "      sigmao2hat_tp1: estimate of prior variance sigma0^2 at round t+1 (for action2).\n",
        "  '''\n",
        "\n",
        "  sigmao1hat_tp1 = (1-1/roundt) * sigmao1hat_t + (1/roundt) * np.abs(P1_avgm_roundt - P1ohat_t) * np.abs(P1_avgm_roundt - P1ohat_tp1)\n",
        "  sigmao2hat_tp1 = (1-1/roundt) * sigmao2hat_t + (1/roundt) * np.abs(P2_avgm_roundt - P2ohat_t) * np.abs(P2_avgm_roundt - P2ohat_tp1)\n",
        "\n",
        "  assert np.all((sigmao1hat_tp1 >=0) & (sigmao2hat_tp1 >=0) )\n",
        "\n",
        "  return sigmao1hat_tp1, sigmao2hat_tp1"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "LrScHEJFwu0C"
      },
      "outputs": [],
      "source": [
        "# @title update_alpha\n",
        "def update_alpha(sigma0_1, sigma0_2, samples_,\n",
        "                 sa_counts_avg, v2, roundt):\n",
        "  '''\n",
        "  alpha_t =  \\frac{v^2}{\\sigma^2_o * m + v^2}\n",
        "  '''\n",
        "\n",
        "  alpha1_t = 0.01*np.ones((sigma0_1.shape[0], 1))\n",
        "  alpha2_t = 0.01*np.ones((sigma0_2.shape[0], 1))\n",
        "\n",
        "  for i in range(sigma0_1.shape[0]):\n",
        "    alpha1_t[i,0] = v2/(np.multiply(np.max(sigma0_1[i, :]), samples_) + v2) \n",
        "    alpha2_t[i,0] = v2/(np.multiply(np.max(sigma0_2[i, :]), samples_) + v2) \n",
        "\n",
        "  assert np.all((alpha1_t >=0) & (alpha1_t <=1) )\n",
        "  assert np.all((alpha2_t >=0) & (alpha2_t <=1) )\n",
        "\n",
        "  return alpha1_t, alpha2_t"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "PoyVTAQlYlNp"
      },
      "outputs": [],
      "source": [
        "# @title update_alpha_new\n",
        "def update_alpha_new(sigma0_1, sigma0_2, samples_,\n",
        "                 sa_counts_avg, v2, roundt):\n",
        "  '''\n",
        "  alpha_t =  \\frac{v^2}{\\sigma^2_o * (1+1/t) * m + v^2}\n",
        "  '''\n",
        "\n",
        "  alpha1_t = 0.01*np.ones((sigma0_1.shape[0], 1))\n",
        "  alpha2_t = 0.01*np.ones((sigma0_2.shape[0], 1))\n",
        "\n",
        "  for i in range(sigma0_1.shape[0]):\n",
        "    alpha1_t[i,0] = v2/(np.multiply(np.multiply(np.max(sigma0_1[i, :]), (1+1/roundt)), samples_) + v2) \n",
        "    alpha2_t[i,0] = v2/(np.multiply(np.multiply(np.max(sigma0_2[i, :]), (1+1/roundt)), samples_) + v2)    \n",
        "\n",
        "  assert np.all((alpha1_t >=0) & (alpha1_t <=1) )\n",
        "  assert np.all((alpha2_t >=0) & (alpha2_t <=1) )\n",
        "\n",
        "  return alpha1_t, alpha2_t"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "eUavJ4q3_G3t"
      },
      "outputs": [],
      "source": [
        "# @title update_alpha_manual\n",
        "def update_alpha_manual(sigma0_1, sigma0_2, samples_,\n",
        "                 sa_counts_avg, v2, roundt):\n",
        "  '''\n",
        "  alpha_t =  \\frac{v^2}{\\sigma^2_o * m + v^2}\n",
        "  '''\n",
        "\n",
        "  alpha1_t = np.ones((sigma0_1.shape[0], 1))\n",
        "  alpha2_t = np.ones((sigma0_1.shape[0], 1))\n",
        "\n",
        "  assert np.all((alpha1_t >=0) & (alpha1_t <=1) )\n",
        "  assert np.all((alpha2_t >=0) & (alpha2_t <=1) )\n",
        "\n",
        "  return alpha1_t, alpha2_t"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "fI_Fo1XHlvW3"
      },
      "source": [
        "# Parallelized Code"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "zqT3Uop-Tjcy"
      },
      "outputs": [],
      "source": [
        "#  @title g3_mp\n",
        "from concurrent import futures\n",
        "from typing import Any, Callable, Iterable, List, Optional, Text, Tuple\n",
        "import cloudpickle\n",
        "import functools\n",
        "\n",
        "def _serialize_closure(closure: Callable) -> Text:\n",
        "  return cloudpickle.dumps(closure)\n",
        "\n",
        "def _run_closure(serialized_closure: bytes) -> Any:\n",
        "  return cloudpickle.loads(serialized_closure)()\n",
        "\n",
        "def pmap(\n",
        "    func: Callable[..., Any],\n",
        "    iterables: Iterable[Any],\n",
        "    nproc: Optional[int] = 0\n",
        ") -> List[Any]:\n",
        "  \"\"\"Maps callable to iterables using multiple processes.\n",
        "\n",
        "  This function is intended to be used like built-in map() function. It is\n",
        "  implemented using g3_multiprocessing. If an exception occurs during the\n",
        "  remote execution, this exception is re-raised in the calling process.\n",
        "\n",
        "  Examples:\n",
        "  \n",
        "  pmap(abs, [-1, 2, -3])\n",
        "  ==> [1, 2, 3]\n",
        "\n",
        "  pmap(lambda x, y: x + y, [1, 2, 3], [4, 5, 6], nproc=2)\n",
        "  ==> [5, 7, 9]\n",
        "\n",
        "  Args:\n",
        "    func: Function to apply to iterables.\n",
        "    iterables: One or more positional args to be passed to `func`.\n",
        "    nproc: use that many Python processes. Value 0 is special, it just calls\n",
        "      built-in map(), which is useful for debugging.\n",
        "\n",
        "  Returns:\n",
        "    list_of_results, equivalent to list(map(func, *iterables)).\n",
        "  \"\"\"\n",
        "\n",
        "  if nproc == 0:\n",
        "    return list(map(func, iterables))\n",
        "\n",
        "  with futures.ProcessPoolExecutor(max_workers=nproc) as executor:\n",
        "    return list(\n",
        "        executor.map(\n",
        "            _run_closure,\n",
        "            (_serialize_closure(functools.partial(func, *arg))\n",
        "             for arg in iterables)))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "AvtCqSuJTugp",
        "outputId": "868dbdf9-582e-4e58-a4c0-15bf01a0e948"
      },
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "[None, None, None, None, None, None, None, None]"
            ]
          },
          "metadata": {},
          "execution_count": 35
        }
      ],
      "source": [
        "# @title example code\n",
        "import time\n",
        "\n",
        "def dummy(s):\n",
        "  time.sleep(s)\n",
        "\n",
        "%timeit\n",
        "pmap(dummy, [[10]]* 8, nproc=8)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "L7XGnmijmRC4"
      },
      "outputs": [],
      "source": [
        "#@title ada-pomrl experiment to parallelize\n",
        "def experiment_omrl(num_tasks, gammas, num_states, k, samples_, trials,\n",
        "                    trueGamma, v2, max_trajectory_length,\n",
        "                    seed_in):\n",
        "  # planning loss per round t\n",
        "  loss_per_round =  np.zeros((num_tasks, len(gammas)))\n",
        "  err_per_round  =  np.zeros((num_tasks, len(gammas)))\n",
        "  optGamma_round = np.zeros((1, num_tasks))\n",
        "  var0_gap_per_round = np.zeros((1, num_tasks))\n",
        "  var1_gap_per_round = np.zeros((1, num_tasks))\n",
        "  mean0_gap_per_round = np.zeros((1, num_tasks))\n",
        "  mean1_gap_per_round = np.zeros((1, num_tasks))\n",
        "\n",
        "\n",
        "  # initialize alpha\n",
        "  alpha1_t = 0.0 * np.ones((num_states,1))\n",
        "  alpha2_t = 0.0 * np.ones((num_states,1))\n",
        "  alpha_t = [alpha1_t, alpha2_t]\n",
        "\n",
        "\n",
        "  #Run specific seed\n",
        "  np.random.seed(seed_in)\n",
        "\n",
        "  # Generate prior distribution from randomMDP\n",
        "  P1_Prior, R1 = randomMDP(num_states, k)\n",
        "  P2_Prior, R2 = randomMDP(num_states, k)\n",
        "\n",
        "  sigma0_1 = np.zeros((num_states, num_states))\n",
        "  sigma0_2 = np.zeros((num_states, num_states))\n",
        "\n",
        "  # define the dirichlet prior distribution\n",
        "  alphasdirich1 = (P1_Prior + 0.015)*1.5*num_states\n",
        "  alphasdirich2 = (P2_Prior + 0.015)*1.5*num_states\n",
        "\n",
        "  # ground truth variance of prior distribution\n",
        "  for i in range(0, num_states):\n",
        "    z = np.sum(alphasdirich1[i,:])\n",
        "    sigma0_1[i,:] = np.multiply(alphasdirich1[i,:]/z, (1-alphasdirich1[i,:]/z))\n",
        "    sigma0_1[i,:] = sigma0_1[i,:]/(z+1)\n",
        "\n",
        "    z = np.sum(alphasdirich2[i,:])\n",
        "    sigma0_2[i,:] = np.multiply(alphasdirich2[i,:]/z, (1-alphasdirich2[i,:]/z))\n",
        "    sigma0_2[i,:] = sigma0_2[i,:]/(z+1)\n",
        "\n",
        "  # set meta-initialize P_o_hat, sigma_o_hat, alpha_t > 0\n",
        "  # uniform prior initialization \n",
        "  P1ohat_t = np.ones((num_states, num_states))/num_states\n",
        "  P2ohat_t = np.ones((num_states, num_states))/num_states\n",
        "\n",
        "  sigmao1hat_t = 0.1 * np.ones((num_states, num_states)) # consider initialization with 0.25\n",
        "  sigmao2hat_t = 0.1 * np.ones((num_states, num_states))\n",
        "\n",
        "  P1ohat_tp1 = np.zeros((num_states, num_states))\n",
        "  P2ohat_tp1 = np.zeros((num_states, num_states))\n",
        "  sigmao1hat_tp1 = 0.1 * np.ones((num_states, num_states)) #np.zeros((num_states, num_states))\n",
        "  sigmao2hat_tp1 = 0.1 * np.ones((num_states, num_states)) #np.zeros((num_states, num_states))\n",
        "\n",
        "  for roundt in range(1, num_tasks+1):\n",
        "      # sample a Pt ~ task dist. for round t\n",
        "      Pt_1 = np.zeros((num_states, num_states))\n",
        "      Pt_2 = np.zeros((num_states, num_states))  \n",
        "      for i in range(0, num_states):\n",
        "        Pt_1[i,:] = np.random.dirichlet(alphasdirich1[i,:])\n",
        "        Pt_2[i,:] = np.random.dirichlet(alphasdirich2[i,:])\n",
        "      _unit_test_P(Pt_1)\n",
        "      _unit_test_P(Pt_2) \n",
        "      \n",
        "      # generate a batch of data and process to compute planning loss\n",
        "      loss_t, err_t = processwithintask(Pt_1, R1, Pt_2, R2, samples_, \n",
        "                                    gammas, trueGamma, trials,\n",
        "                                    alpha=alpha_t, P1o=P1ohat_t,\n",
        "                                    P2o=P2ohat_t)\n",
        "                                    \n",
        "      idx = np.argmin(loss_t)\n",
        "      optGamma_round[0, roundt-1] = gammas[idx]\n",
        "\n",
        "      loss_per_round[roundt-1, :] = loss_t\n",
        "      err_per_round[roundt-1, :] = err_t\n",
        "      \n",
        "      # online average of estimates - meta learner\n",
        "      P1ohat_tp1, P2ohat_tp1, P1_avgm_roundt, P2_avgm_roundt, sa_counts_avg = \\\n",
        "                                              update_prior_online(P1ohat_t, P2ohat_t,\n",
        "                                                                  roundt, samples_, Pt_1, Pt_2, R1, R2,\n",
        "                                                                  max_trajectory_length, trials)\n",
        "\n",
        "\n",
        "      alpha1_tp1, alpha2_tp1 = update_alpha_new(sigmao1hat_tp1, sigmao2hat_tp1,\n",
        "                                            samples_, sa_counts_avg, v2,\n",
        "                                            roundt)  \n",
        "      \n",
        "      alpha1_t = alpha1_tp1\n",
        "      alpha2_t = alpha2_tp1\n",
        "      alpha_t = [alpha1_t, alpha2_t]\n",
        "      \n",
        "      \n",
        "\n",
        "      if roundt==1:\n",
        "        continue\n",
        "      \n",
        "      # compute meta-updates\n",
        "      sigmao1hat_tp1, sigmao2hat_tp1 = update_variance_online(sigmao1hat_t, sigmao2hat_t,\n",
        "                                                              P1ohat_t, P2ohat_t,\n",
        "                                                              P1ohat_tp1, P2ohat_tp1,\n",
        "                                                              P1_avgm_roundt, P2_avgm_roundt,\n",
        "                                                              roundt)\n",
        "      \n",
        "      \n",
        "      P1ohat_t = P1ohat_tp1\n",
        "      P2ohat_t = P2ohat_tp1\n",
        "\n",
        "      sigmao1hat_t = sigmao1hat_tp1\n",
        "      sigmao2hat_t = sigmao2hat_tp1\n",
        "\n",
        "      var0_gap_per_round[0, roundt-1] = np.linalg.norm(sigma0_1-sigmao1hat_t, ord=2)\n",
        "      var1_gap_per_round[0, roundt-1] = np.linalg.norm(sigma0_2-sigmao2hat_t, ord=2)\n",
        "      mean0_gap_per_round[0, roundt-1] = np.linalg.norm(P1_Prior-P1ohat_t, ord=2)\n",
        "      mean1_gap_per_round[0, roundt-1] = np.linalg.norm(P2_Prior-P2ohat_t, ord=2)\n",
        "\n",
        "  return loss_per_round, err_per_round, optGamma_round, var0_gap_per_round, var1_gap_per_round, mean0_gap_per_round, mean1_gap_per_round"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "YaN-z8LcyYzR"
      },
      "outputs": [],
      "source": [
        "#@title aggregate Baseline to parallelize\n",
        "def experiment_omrl_naive(num_tasks, gammas, num_states, k, samples_, trials,\n",
        "                    trueGamma, v2, max_trajectory_length,\n",
        "                    seed_in):\n",
        "  # planning loss per round t\n",
        "  loss_per_round =  np.zeros((num_tasks, len(gammas)))\n",
        "  err_per_round  =  np.zeros((num_tasks, len(gammas)))\n",
        "  optGamma_round = np.zeros((1, num_tasks))\n",
        "  var0_gap_per_round = np.zeros((1, num_tasks))\n",
        "  var1_gap_per_round = np.zeros((1, num_tasks))\n",
        "  mean0_gap_per_round = np.zeros((1, num_tasks))\n",
        "  mean1_gap_per_round = np.zeros((1, num_tasks))\n",
        "\n",
        "\n",
        "  # initialize alpha\n",
        "  alpha1_t = np.ones((num_states,1))\n",
        "  alpha2_t = np.ones((num_states,1))\n",
        "  alpha_t = [alpha1_t, alpha2_t]\n",
        "\n",
        "\n",
        "  #Run specific seed\n",
        "  np.random.seed(seed_in)\n",
        "\n",
        "  # Generate prior distribution from randomMDP\n",
        "  P1_Prior, R1 = randomMDP(num_states, k)\n",
        "  P2_Prior, R2 = randomMDP(num_states, k)\n",
        "\n",
        "  sigma0_1 = np.zeros((num_states, num_states))\n",
        "  sigma0_2 = np.zeros((num_states, num_states))\n",
        "\n",
        "  # define the dirichlet prior distribution\n",
        "  alphasdirich1 = (P1_Prior + 0.015)*1.5*num_states\n",
        "  alphasdirich2 = (P2_Prior + 0.015)*1.5*num_states\n",
        "\n",
        "  # ground truth variance of prior distribution\n",
        "  for i in range(0, num_states):\n",
        "    z = np.sum(alphasdirich1[i,:])\n",
        "    sigma0_1[i,:] = np.multiply(alphasdirich1[i,:]/z, (1-alphasdirich1[i,:]/z))\n",
        "    sigma0_1[i,:] = sigma0_1[i,:]/(z+1)\n",
        "\n",
        "    z = np.sum(alphasdirich2[i,:])\n",
        "    sigma0_2[i,:] = np.multiply(alphasdirich2[i,:]/z, (1-alphasdirich2[i,:]/z))\n",
        "    sigma0_2[i,:] = sigma0_2[i,:]/(z+1)\n",
        "\n",
        "\n",
        "  # set meta-initialize P_o_hat, sigma_o_hat, alpha_t > 0\n",
        "  # uniform prior initialization \n",
        "  P1ohat_t = np.ones((num_states, num_states))/num_states\n",
        "  P2ohat_t = np.ones((num_states, num_states))/num_states\n",
        "\n",
        "  sigmao1hat_t = 0.1 * np.ones((num_states, num_states)) # consider initialization with 0.25\n",
        "  sigmao2hat_t = 0.1 * np.ones((num_states, num_states))\n",
        "\n",
        "  P1ohat_tp1 = np.zeros((num_states, num_states))\n",
        "  P2ohat_tp1 = np.zeros((num_states, num_states))\n",
        "  sigmao1hat_tp1 = 0.1 * np.ones((num_states, num_states)) #np.zeros((num_states, num_states))\n",
        "  sigmao2hat_tp1 = 0.1 * np.ones((num_states, num_states)) #np.zeros((num_states, num_states))\n",
        "\n",
        "  for roundt in range(1, num_tasks+1):\n",
        "      # sample a Pt ~ task dist. for round t\n",
        "      Pt_1 = np.zeros((num_states, num_states))\n",
        "      Pt_2 = np.zeros((num_states, num_states))  \n",
        "      for i in range(0, num_states):\n",
        "        Pt_1[i,:] = np.random.dirichlet(alphasdirich1[i,:])\n",
        "        Pt_2[i,:] = np.random.dirichlet(alphasdirich2[i,:])\n",
        "      _unit_test_P(Pt_1)\n",
        "      _unit_test_P(Pt_2) \n",
        "      \n",
        "      # generate a batch of data and process to compute planning loss\n",
        "      # For t=1, naive baseline should only consider samples too.\n",
        "      if roundt==1:\n",
        "        alpha1_t = 0.0 * np.ones((num_states,1))\n",
        "        alpha2_t = 0.0 * np.ones((num_states,1))\n",
        "      else:\n",
        "        alpha1_t = np.ones((num_states,1))\n",
        "        alpha2_t = np.ones((num_states,1))\n",
        "             \n",
        "                     \n",
        "      alpha_t = [alpha1_t, alpha2_t]\n",
        "      loss_t, err_t = processwithintask(Pt_1, R1, Pt_2, R2, samples_, \n",
        "                                    gammas, trueGamma, trials,\n",
        "                                    alpha=alpha_t, P1o=P1ohat_t,\n",
        "                                    P2o=P2ohat_t)\n",
        "                                    \n",
        "      idx = np.argmin(loss_t)\n",
        "      optGamma_round[0, roundt-1] = gammas[idx]\n",
        "\n",
        "      loss_per_round[roundt-1, :] = loss_t\n",
        "      err_per_round[roundt-1, :] = err_t\n",
        "      \n",
        "      # online average of estimates - meta learner\n",
        "      P1ohat_tp1, P2ohat_tp1, P1_avgm_roundt, P2_avgm_roundt, sa_counts_avg = \\\n",
        "                                              update_prior_online(P1ohat_t, P2ohat_t,\n",
        "                                                                  roundt, samples_, Pt_1, Pt_2, R1, R2,\n",
        "                                                                  max_trajectory_length, trials)\n",
        "      \n",
        "      if roundt==1:\n",
        "        continue\n",
        "      \n",
        "\n",
        "      P1ohat_t = P1ohat_tp1\n",
        "      P2ohat_t = P2ohat_tp1\n",
        "\n",
        "      var0_gap_per_round[0, roundt-1] = np.linalg.norm(sigma0_1-sigmao1hat_t, ord=2)\n",
        "      var1_gap_per_round[0, roundt-1] = np.linalg.norm(sigma0_2-sigmao2hat_t, ord=2)\n",
        "      mean0_gap_per_round[0, roundt-1] = np.linalg.norm(P1_Prior-P1ohat_t, ord=2)\n",
        "      mean1_gap_per_round[0, roundt-1] = np.linalg.norm(P2_Prior-P2ohat_t, ord=2)\n",
        "\n",
        "  return loss_per_round, err_per_round, optGamma_round, var0_gap_per_round, var1_gap_per_round, mean0_gap_per_round, mean1_gap_per_round"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "lMFSV4GS-8zM"
      },
      "outputs": [],
      "source": [
        "#@title pomrl(known sigma) experiment to parallelize\n",
        "def experiment_omrl_knownsigma(num_tasks, gammas, num_states, k, samples_, trials,\n",
        "                    trueGamma, v2, max_trajectory_length,\n",
        "                    seed_in):\n",
        "  # planning loss per round t\n",
        "  loss_per_round =  np.zeros((num_tasks, len(gammas)))\n",
        "  err_per_round  =  np.zeros((num_tasks, len(gammas)))\n",
        "  optGamma_round = np.zeros((1, num_tasks))\n",
        "  var0_gap_per_round = np.zeros((1, num_tasks))\n",
        "  var1_gap_per_round = np.zeros((1, num_tasks))\n",
        "  mean0_gap_per_round = np.zeros((1, num_tasks))\n",
        "  mean1_gap_per_round = np.zeros((1, num_tasks))\n",
        "\n",
        "\n",
        "  # initialize alpha\n",
        "  alpha1_t = 0.0 * np.ones((num_states,1))\n",
        "  alpha2_t = 0.0 * np.ones((num_states,1))\n",
        "  alpha_t = [alpha1_t, alpha2_t]\n",
        "\n",
        "\n",
        "  #Run specific seed\n",
        "  np.random.seed(seed_in)\n",
        "\n",
        "  # Generate prior distribution from randomMDP\n",
        "  P1_Prior, R1 = randomMDP(num_states, k)\n",
        "  P2_Prior, R2 = randomMDP(num_states, k)\n",
        "\n",
        "  sigma0_1 = np.zeros((num_states, num_states))\n",
        "  sigma0_2 = np.zeros((num_states, num_states))\n",
        "\n",
        "  # define the dirichlet prior distribution\n",
        "  alphasdirich1 = (P1_Prior + 0.015)*1.5*num_states\n",
        "  alphasdirich2 = (P2_Prior + 0.015)*1.5*num_states\n",
        "\n",
        "  # ground truth variance of prior distribution\n",
        "  for i in range(0, num_states):\n",
        "    z = np.sum(alphasdirich1[i,:])\n",
        "    sigma0_1[i,:] = np.multiply(alphasdirich1[i,:]/z, (1-alphasdirich1[i,:]/z))\n",
        "    sigma0_1[i,:] = sigma0_1[i,:]/(z+1)\n",
        "\n",
        "    z = np.sum(alphasdirich2[i,:])\n",
        "    sigma0_2[i,:] = np.multiply(alphasdirich2[i,:]/z, (1-alphasdirich2[i,:]/z))\n",
        "    sigma0_2[i,:] = sigma0_2[i,:]/(z+1)\n",
        "\n",
        "\n",
        "  # set meta-initialize P_o_hat, sigma_o_hat, alpha_t > 0\n",
        "  # uniform prior initialization \n",
        "  P1ohat_t = np.ones((num_states, num_states))/num_states\n",
        "  P2ohat_t = np.ones((num_states, num_states))/num_states\n",
        "\n",
        "  sigmao1hat_t = 0.1 * np.ones((num_states, num_states)) # consider initialization with 0.25\n",
        "  sigmao2hat_t = 0.1 * np.ones((num_states, num_states))\n",
        "\n",
        "  P1ohat_tp1 = np.zeros((num_states, num_states))\n",
        "  P2ohat_tp1 = np.zeros((num_states, num_states))\n",
        "  sigmao1hat_tp1 = 0.1 * np.ones((num_states, num_states)) #np.zeros((num_states, num_states))\n",
        "  sigmao2hat_tp1 = 0.1 * np.ones((num_states, num_states)) #np.zeros((num_states, num_states))\n",
        "\n",
        "  for roundt in range(1, num_tasks+1):\n",
        "      # sample a Pt ~ task dist. for round t\n",
        "      Pt_1 = np.zeros((num_states, num_states))\n",
        "      Pt_2 = np.zeros((num_states, num_states))  \n",
        "      for i in range(0, num_states):\n",
        "        Pt_1[i,:] = np.random.dirichlet(alphasdirich1[i,:])\n",
        "        Pt_2[i,:] = np.random.dirichlet(alphasdirich2[i,:])\n",
        "      _unit_test_P(Pt_1)\n",
        "      _unit_test_P(Pt_2) \n",
        "      \n",
        "      # generate a batch of data and process to compute planning loss\n",
        "      loss_t, err_t = processwithintask(Pt_1, R1, Pt_2, R2, samples_, \n",
        "                                    gammas, trueGamma, trials,\n",
        "                                    alpha=alpha_t, P1o=P1ohat_t,\n",
        "                                    P2o=P2ohat_t)\n",
        "                                    \n",
        "      idx = np.argmin(loss_t)\n",
        "      optGamma_round[0, roundt-1] = gammas[idx]\n",
        "\n",
        "      loss_per_round[roundt-1, :] = loss_t\n",
        "      err_per_round[roundt-1, :] = err_t\n",
        "      \n",
        "      # online average of estimates - meta learner\n",
        "      P1ohat_tp1, P2ohat_tp1, P1_avgm_roundt, P2_avgm_roundt, sa_counts_avg = \\\n",
        "                                              update_prior_online(P1ohat_t, P2ohat_t,\n",
        "                                                                  roundt, samples_, Pt_1, Pt_2, R1, R2,\n",
        "                                                                  max_trajectory_length, trials)\n",
        "\n",
        "\n",
        "      alpha1_tp1, alpha2_tp1 = update_alpha_new(sigma0_1, sigma0_2,\n",
        "                                            samples_, sa_counts_avg, v2,\n",
        "                                            roundt)  \n",
        " \n",
        "      alpha1_t = alpha1_tp1\n",
        "      alpha2_t = alpha2_tp1\n",
        "      alpha_t = [alpha1_t, alpha2_t]\n",
        "      \n",
        "      \n",
        "      if roundt==1:\n",
        "        continue\n",
        "      \n",
        "      # compute meta-updates\n",
        "      sigmao1hat_tp1, sigmao2hat_tp1 = sigma0_1, sigma0_2\n",
        "      \n",
        "      \n",
        "      P1ohat_t = P1ohat_tp1\n",
        "      P2ohat_t = P2ohat_tp1\n",
        "\n",
        "      sigmao1hat_t = sigmao1hat_tp1\n",
        "      sigmao2hat_t = sigmao2hat_tp1\n",
        "\n",
        "      var0_gap_per_round[0, roundt-1] = np.linalg.norm(sigma0_1-sigmao1hat_t, ord=2)\n",
        "      var1_gap_per_round[0, roundt-1] = np.linalg.norm(sigma0_2-sigmao2hat_t, ord=2)\n",
        "      mean0_gap_per_round[0, roundt-1] = np.linalg.norm(P1_Prior-P1ohat_t, ord=2)\n",
        "      mean1_gap_per_round[0, roundt-1] = np.linalg.norm(P2_Prior-P2ohat_t, ord=2)\n",
        "\n",
        "  return loss_per_round, err_per_round, optGamma_round, var0_gap_per_round, var1_gap_per_round, mean0_gap_per_round, mean1_gap_per_round"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Main Script"
      ],
      "metadata": {
        "id": "lHKSWMNGsDk9"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "JpBVuasCl2ta",
        "cellView": "form"
      },
      "outputs": [],
      "source": [
        "# @title call to experiments\n",
        "seed = 88 \n",
        "samples_ = 5\n",
        "num_tasks = 15\n",
        "num_states = 10\n",
        "max_trajectory_length = 5\n",
        "k = 5\n",
        "v2 = 1/4\n",
        "gammas = [gamma_ for gamma_ in np.arange(0, 0.99, 0.1)]\n",
        "gammas.append(0.99)\n",
        "trueGamma = max(gammas)\n",
        "runs = 50\n",
        "trials = 100\n",
        "\n",
        "iterables = []\n",
        "for run_id in range(runs):\n",
        "  iterables.append([num_tasks, gammas, num_states, k, samples_, \n",
        "                            trials, trueGamma, v2, max_trajectory_length, seed+run_id])\n",
        "\n",
        "%timeit\n",
        "\n",
        "# runs ada-pomrl baseline\n",
        "A = pmap(experiment_omrl, iterables, nproc=8) \n",
        "\n",
        "# runs aggregation baseline\n",
        "# A = pmap(experiment_omrl_naive, iterables, nproc=8)\n",
        "\n",
        "# runs pomrl- known sigma baseline\n",
        "# A = pmap(experiment_omrl_knownsigma, iterables, nproc=8)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "cJk1qNrdWMxN"
      },
      "source": [
        "# Data Logging"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "OTPraU4oFkIe"
      },
      "source": [
        "## Write data to grdive"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "mT9s955eFmBm",
        "outputId": "dafcadae-cbfd-4b3f-9d99-ccc1f18a2191"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Mounted at /content/drive\n"
          ]
        }
      ],
      "source": [
        "import os\n",
        "import numpy as np\n",
        "from numpy import save\n",
        "from numpy import load\n",
        "from google.colab import files\n",
        "from google.colab import drive\n",
        "drive.mount('/content/drive', force_remount=True)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "EoqqEenjF2_p"
      },
      "outputs": [],
      "source": [
        "filename = 'WithoutMetaLearning_2022_sacounts_exp_dirich_0PT015_1PT5N_50runs_15tasks_5samples_5trajle_gammaPT90.npy'\n",
        "fullpath = '/content/drive/My Drive/onlinemetarl_experiments/' + filename"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "qbMfku2GJwPl"
      },
      "outputs": [],
      "source": [
        "np.save(fullpath, A)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "zkEPqVzlNU7A"
      },
      "outputs": [],
      "source": [
        "os.chdir(\"/content/drive/My Drive/onlinemetarl_experiments\")\n",
        "!ls"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "2dqOproZTrNF"
      },
      "outputs": [],
      "source": [
        "# @title Encode with EnvLogger's codec\n",
        "encoded_z = codec.encode(A)\n",
        "with open(fullpath, 'wb') as f:\n",
        "  f.write(encoded_z.SerializeToString())"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "FzR7ZSAkUEBy"
      },
      "outputs": [],
      "source": [
        "# @title Decode with EnvLogger's codec\n",
        "with open(fullpath, 'rb') as f:\n",
        "  proto = storage_pb2.Data()\n",
        "  proto.ParseFromString(f.read())\n",
        "  decoded_A = codec.decode(proto)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3G3q3UouX4HC"
      },
      "source": [
        "# Selecion of $\\gamma$"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "DHdEfnmifyt6"
      },
      "source": [
        "## Utilities"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Z1xM7LV1rdTh",
        "cellView": "form"
      },
      "outputs": [],
      "source": [
        "# @title Load Saved Data\n",
        "\n",
        "num_tasks = 15 # @param{\"type\": \"string\"}\n",
        "runs = 100 # @param{\"type\": \"string\"}\n",
        "filename = 'Ours_2022_sacounts_exp_dirich_0PT015BY4_1PT5N_100runs_15tasks_5samples_5trajle.npy' # @param{\"type\": \"string\"}\n",
        "DIR ='/content/drive/My Drive/onlinemetarl_experiments'  # @param{\"type\": \"string\"}\n",
        "\n",
        "gammas = [gamma_ for gamma_ in np.arange(0, 0.99, 0.1)]\n",
        "gammas.append(0.99)\n",
        "trueGamma = max(gammas)\n",
        "\n",
        "\n",
        "fullpath = os.path.join(DIR, filename)\n",
        "with open(fullpath, 'rb') as f:\n",
        "  proto = storage_pb2.Data()\n",
        "  proto.ParseFromString(f.read())\n",
        "  decoded_A = codec.decode(proto)\n",
        "\n",
        "loss_per_run =  np.zeros((runs, num_tasks, len(gammas)))\n",
        "err_per_run  =  np.zeros((runs, num_tasks, len(gammas)))\n",
        "optGamma_run = np.zeros((runs, num_tasks))\n",
        "\n",
        "avg_loss_per_round = np.zeros((num_tasks, len(gammas)))\n",
        "avg_err_per_round = np.zeros((num_tasks, len(gammas)))\n",
        "\n",
        "for i in range(runs):\n",
        "  loss_per_run[i, :, :] = decoded_A[i][0]\n",
        "  err_per_run[i, :, :] = decoded_A[i][1]\n",
        "  optGamma_run[i,:] = decoded_A[i][2]\n",
        "\n",
        "avg_loss_per_round = np.mean(loss_per_run, axis=0)\n",
        "avg_err_per_round = np.mean(err_per_run, axis=0)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "gaTe-Yd9h1h1"
      },
      "outputs": [],
      "source": [
        "# @title optgamma_ourproposal_alpha\n",
        "def update_alpha_new_uni(sigma0, samples_, roundt):\n",
        "  '''\n",
        "  alpha_t =  \\frac{v^2}{\\sigma^2_o * (1+1/t) * m + 1}\n",
        "  '''\n",
        "  sigma0_square = np.square(sigma0)\n",
        "  alpha_t = 1/(np.multiply(np.multiply(sigma0_square,(1+1/roundt)), samples_) + 1)   \n",
        "  return alpha_t\n",
        "\n",
        "def compute_optimal_alpha_uni(sigma0, samples_):\n",
        "  alpha_star = 1/(np.multiply(sigma0, samples_) + 1)\n",
        "  return alpha_star\n",
        "\n",
        "def optgamma_ourproposal_alpha(numtasks, m, \n",
        "                               sigma_true,\n",
        "                               schedule=1/5, \n",
        "                               tau=None):  \n",
        "\n",
        "  def tautogamma(tau):\n",
        "    gamma = 1-1/tau\n",
        "    return gamma\n",
        "\n",
        "  sigma_init = 0.25\n",
        "  sigma_true = sigma_true\n",
        "  sigma_list = np.arange(sigma_init, sigma_true, -0.018)\n",
        "  sigma_t = [i for i in sigma_list]\n",
        "  sigma_t.append(sigma_t[-1])\n",
        "  sigma_t.append(sigma_t[-1])\n",
        "  \n",
        "\n",
        "  optGamma = np.zeros((1, numtasks))\n",
        "  To = m*2 #t=1\n",
        "  tau = pow(To, schedule)\n",
        "  optGamma[0,0] = tautogamma(tau)\n",
        "  \n",
        "  for roundt in range(1, numtasks):\n",
        "    alpha_t = update_alpha_new_uni(sigma_t[roundt], m, roundt)\n",
        "    Tkt = np.multiply((1 - alpha_t), m) + np.multiply(alpha_t, m*(roundt-1))\n",
        "    tau = pow(Tkt*2, schedule)\n",
        "    optGamma[0,roundt] = tautogamma(tau)\n",
        "  return optGamma"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "AB0MRWqdih-H"
      },
      "outputs": [],
      "source": [
        "# @title optgamma_analytical_fn\n",
        "def optgamma_analytical_fn(numtasks, m, num_states, sigma, gamma_0):  \n",
        "  '''\n",
        "  gamma  = min(1, gamma_0+ gamma_tilde)\n",
        "\n",
        "  '''\n",
        "  def compute_gamma_tilde(C):\n",
        "    if C >= 1 :\n",
        "      return 0\n",
        "    elif C < 1/2 :\n",
        "      return 1\n",
        "    else:\n",
        "      answer = (1-C)/(1+C)\n",
        "      return answer\n",
        "\n",
        "  def compute_C(sigma, m, T, num_states):\n",
        "    sigma2m_1 = np.multiply(np.square(sigma), m) + 1\n",
        "    C1 = 2/np.multiply(sigma, m)\n",
        "    C2 = ((1/np.sqrt( T+ 1)) * (sigma + np.sqrt(1/m))) / (sigma2m_1)\n",
        "    C3 = np.multiply(sigma2m_1, np.sqrt(1/m)) / (sigma2m_1)\n",
        "    C = (C2 + C3)\n",
        "    return C\n",
        "\n",
        "  optGamma = np.zeros((1, numtasks))  \n",
        "  for roundt in range(0, numtasks):\n",
        "    C = compute_C(sigma, m, roundt, num_states)\n",
        "    gamma_tilde = compute_gamma_tilde(C)\n",
        "    gamma_temp = gamma_0+gamma_tilde\n",
        "    optGamma[0,roundt] = min(1, gamma_temp)\n",
        "  return optGamma"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "OcmMEk_sbFYS"
      },
      "source": [
        "## Compile Results"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "iCiJD5zzmDLc"
      },
      "outputs": [],
      "source": [
        "# @title Load Saved Data\n",
        "\n",
        "\n",
        "num_tasks = 15 # @param{\"type\": \"string\"}\n",
        "runs = 100 # @param{\"type\": \"string\"}\n",
        "filename = 'Ours_2022_sacounts_exp_dirich_0PT015BY4_1PT5N_100runs_15tasks_5samples_5trajle.npy' # @param{\"type\": \"string\"}\n",
        "DIR ='/content/drive/My Drive/onlinemetarl_experiments'  # @param{\"type\": \"string\"}\n",
        "\n",
        "gammas = [gamma_ for gamma_ in np.arange(0, 0.99, 0.1)]\n",
        "gammas.append(0.99)\n",
        "trueGamma = max(gammas)\n",
        "\n",
        "\n",
        "fullpath = os.path.join(DIR, filename)\n",
        "with open(fullpath, 'rb') as f:\n",
        "  proto = storage_pb2.Data()\n",
        "  proto.ParseFromString(f.read())\n",
        "  decoded_A = codec.decode(proto)\n",
        "\n",
        "loss_per_run =  np.zeros((runs, num_tasks, len(gammas)))\n",
        "err_per_run  =  np.zeros((runs, num_tasks, len(gammas)))\n",
        "optGamma_run = np.zeros((runs, num_tasks))\n",
        "\n",
        "avg_loss_per_round = np.zeros((num_tasks, len(gammas)))\n",
        "avg_err_per_round = np.zeros((num_tasks, len(gammas)))\n",
        "\n",
        "for i in range(runs):\n",
        "  loss_per_run[i, :, :] = decoded_A[i][0]\n",
        "  err_per_run[i, :, :] = decoded_A[i][1]\n",
        "  optGamma_run[i,:] = decoded_A[i][2]\n",
        "\n",
        "avg_loss_per_round = np.mean(loss_per_run, axis=0)\n",
        "avg_err_per_round = np.mean(err_per_run, axis=0)\n",
        "std_dev_per_round = np.std(err_per_run, axis=0) # standard deviation across runs.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "4tyoIu9496Oh",
        "outputId": "c921eabf-2274-4fb0-ea18-39d8745ac8eb"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "0.949638347757141 5\n",
            "0.9926812590884482 6\n",
            "1.0690266251983802 8\n",
            "0.8325349898365892 7\n",
            "0.8377454891372453 8\n",
            "0.9241748117251446 9\n"
          ]
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "<Figure size 640x480 with 1 Axes>"
            ],
            "image/png": "[encoded data, 57064 characters]\n"
          },
          "metadata": {}
        }
      ],
      "source": [
        "# @title Paper OMRL Plots\n",
        "\n",
        "import matplotlib as mpl\n",
        "from matplotlib.ticker import FormatStrFormatter\n",
        "\n",
        "sn.set_context(\"paper\")\n",
        "sn.set_context(\"paper\", font_scale=2.00)\n",
        "plt.rcParams['savefig.facecolor'] = \"0.8\"\n",
        "plt.rcParams['savefig.edgecolor'] = \"black\"\n",
        "plt.rcParams['pdf.fonttype'] = 42\n",
        "plt.rcParams['ps.fonttype'] = 42\n",
        "\n",
        "nlines=8\n",
        "c = np.arange(1, nlines + 1)\n",
        "norm = mpl.colors.Normalize(vmin=c.min()-1, vmax=c.max()+1)\n",
        "cmap = mpl.cm.ScalarMappable(norm=norm, cmap=mpl.cm.bwr)\n",
        "cmap.set_array([])\n",
        "\n",
        "\n",
        "def plotloss_paper(y, e, gammas, numtasks):\n",
        "  fig, ax1 = plt.subplots()\n",
        "  idx = 0\n",
        "  gammalist = []\n",
        "  for i in range(0, numtasks):\n",
        "    if i==0 or i==2 or i==4 or i==5 or i==9 or i==13:\n",
        "      yk = y[i,:].reshape((1,y.shape[-1]))\n",
        "      ek = e[i,:].reshape((1,y.shape[-1]))\n",
        "      plt.errorbar(gammas, yk[0,:], yerr=(ek[0,:]/(np.sqrt(runs))),\n",
        "                  label = 'T={}'.format(int(i+1)),\n",
        "                  color=cmap.to_rgba(idx + i),\n",
        "                  linewidth=2.5)\n",
        "      min_x = np.argmin(yk[0,:])\n",
        "      min_y = np.min(yk[0,:])\n",
        "      print(min_y, min_x)\n",
        "      plt.scatter(gammas[min_x], \n",
        "                  min_y, \n",
        "                  color=cmap.to_rgba(idx + i), \n",
        "                  marker='*',\n",
        "                  s=380)\n",
        "      plt.ylim((0, 4))   # set the ylim to bottom, top\n",
        "      plt.xlim((gammas[0], 1))   # set the ylim to bottom, top\n",
        "      idx += 1\n",
        "      gammalist.append(gammas[min_x])\n",
        "  ax1.set_xlabel('$\\gamma$')\n",
        "  ax1.set_ylabel('Planning Loss')\n",
        "  ax1.set_xticks(gammas)\n",
        "  plt.legend(bbox_to_anchor=(1, 1), \n",
        "             fancybox=True, shadow=True, ncol=4,\n",
        "             facecolor='w', fontsize=16)\n",
        "          \n",
        "  matplotlib.rc('axes', edgecolor='black')\n",
        "  ax1.xaxis.set_major_formatter(FormatStrFormatter('%.1f'))\n",
        "\n",
        "  # plt.savefig((os.path.join(DIR,\"S30m5T20_runs20_randomMDP_averageplanloss_ours_gamma_markers_pruned_1STD\"+ \".pdf\")), \n",
        "  #             dpi=200, orientation='portrait', \n",
        "  #             papertype=None, \n",
        "  #             format=None,\n",
        "  #             transparent=True, \n",
        "  #             bbox_inches='tight', \n",
        "  #             pad_inches=0,\n",
        "  #             figsize=(4, 4))\n",
        "  plt.show()\n",
        "  return fig, gammalist\n",
        "\n",
        "# planning loss per round t\n",
        "loss_per_run =  np.zeros((runs, num_tasks, len(gammas)))\n",
        "err_per_run  =  np.zeros((runs, num_tasks, len(gammas)))\n",
        "optGamma_run = np.zeros((runs, num_tasks))\n",
        "\n",
        "\n",
        "# track estimates per round t\n",
        "var0_gap_run =  np.zeros((runs, num_tasks))\n",
        "var1_gap_run = np.zeros((runs, num_tasks))\n",
        "mean0_gap_run =  np.zeros((runs, num_tasks))\n",
        "mean1_gap_run =  np.zeros((runs, num_tasks))\n",
        "\n",
        "for i in range(runs):\n",
        "  loss_per_run[i, :, :] = decoded_A[i][0]\n",
        "  err_per_run[i, :, :] = decoded_A[i][1]\n",
        "  optGamma_run[i,:] = decoded_A[i][2] \n",
        "  var0_gap_run[i,:] = decoded_A[i][3]\n",
        "  var1_gap_run[i,:] = decoded_A[i][4]\n",
        "  mean0_gap_run[i,:] = decoded_A[i][5]\n",
        "  mean1_gap_run[i,:] = decoded_A[i][6]\n",
        "\n",
        "avg_loss_per_round = np.mean(loss_per_run, axis=0)\n",
        "avg_err_per_round = np.mean(err_per_run, axis=0)\n",
        "std_dev_per_round = np.std(loss_per_run, axis=0)\n",
        "avg_var0_gap = np.mean(var0_gap_run, axis=0)\n",
        "avg_var1_gap = np.mean(var1_gap_run, axis=0)\n",
        "avg_mean0_gap = np.mean(mean0_gap_run, axis=0)\n",
        "avg_mean1_gap = np.mean(mean1_gap_run, axis=0)\n",
        "\n",
        "# 1 STD Dev across runs in Errorbars\n",
        "figomrl_loss, chosen_discount = plotloss_paper(avg_loss_per_round, \n",
        "                                               std_dev_per_round,\n",
        "                                               gammas, num_tasks)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "h_5JIfArynxM",
        "outputId": "d430919c-31d5-4915-8fba-634ed14988b8"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "[0.5, 0.7000000000000001, 0.6000000000000001, 0.6000000000000001, 0.8, 0.7000000000000001, 0.7000000000000001, 0.8, 0.8, 0.8, 0.9, 0.9, 0.9, 0.9, 0.8]\n"
          ]
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "<Figure size 640x480 with 2 Axes>"
            ],
            "image/png": "[encoded data, 84964 characters]\n"
          },
          "metadata": {}
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "[0.5, 0.7000000000000001, 0.6000000000000001, 0.6000000000000001, 0.8, 0.7000000000000001, 0.7000000000000001, 0.8, 0.8, 0.8, 0.9, 0.9, 0.9, 0.9, 0.8]\n"
          ]
        }
      ],
      "source": [
        "#@title Paper Discount Plots\n",
        "\n",
        "sn.set_context(\"paper\")\n",
        "sn.set_context(\"paper\", font_scale=2.00)\n",
        "plt.rcParams['savefig.facecolor'] = \"0.8\"\n",
        "plt.rcParams['savefig.edgecolor'] = \"black\"\n",
        "plt.rcParams['pdf.fonttype'] = 42\n",
        "plt.rcParams['ps.fonttype'] = 42\n",
        "\n",
        "\n",
        "def plotdiscount_paper(y, e, gammas, numtasks): #15*11\n",
        "  xtoplot = list(range(1, num_tasks+1))\n",
        "  fig, ax1 = plt.subplots()\n",
        "  idx = 0\n",
        "  gammalist = []\n",
        "  losslist = []\n",
        "  templist = []\n",
        "  plt.ylim((0.3, 1.4))   # set the ylim to bottom, top\n",
        "  #plt.xlim((1, numtasks+1))   # set the ylim to bottom, top\n",
        "  for i in range(0, numtasks):\n",
        "    if i>=0:\n",
        "      yk = y[i,:].reshape((1,y.shape[-1])) #1*11\n",
        "      ek = e[i,:].reshape((1,y.shape[-1])) #1*11\n",
        "      min_x = np.argmin(yk[0,:])\n",
        "      min_y = np.min(yk[0,:])\n",
        "      plt.scatter(i+1, yk[0,min_x],\n",
        "                  color=cmap.to_rgba(idx + i),\n",
        "                  marker='*',\n",
        "                  s=100)\n",
        "\n",
        "      # use this for 1-std across runs\n",
        "      plt.errorbar(i+1, yk[0,min_x], yerr=(ek[0,min_x]/np.sqrt(runs)),\n",
        "                  color=cmap.to_rgba(idx + i),\n",
        "                   linewidth=2.5)\n",
        "      # plt.errorbar(i+1, yk[0,min_x], yerr=ek[0,min_x],\n",
        "      #             color=cmap.to_rgba(idx + i),\n",
        "      #              linewidth=2.5)\n",
        "\n",
        "      idx += 1\n",
        "      gammalist.append(gammas[min_x])\n",
        "      losslist.append(min_y)\n",
        "      templist.append(yk[0,min_x])\n",
        "\n",
        "  #l1, = plt.plot(xtoplot, gammalist)\n",
        "  plt.plot(xtoplot, templist, color='black', linewidth=2.5)\n",
        "  ax1.set_xticks(xtoplot)\n",
        "  ax1.set_xlabel('Task $T$')\n",
        "  ax1.set_ylabel('Planning Loss', color ='black')\n",
        "  \n",
        "  print(gammalist)\n",
        "  matplotlib.rc('axes', edgecolor='black')\n",
        "  # ax1.xaxis.set_major_formatter(FormatStrFormatter('%.1f'))\n",
        "\n",
        "  ax2 = ax1.twinx()\n",
        "  ax2.set_ylim(0,1)\n",
        "  l2, = ax2.plot(xtoplot, gammalist, \n",
        "                 color=\"tab:green\", \n",
        "                 marker='*', \n",
        "                 markersize='15',\n",
        "                 linewidth=2.5)\n",
        "  ax2.set_ylabel('Empircally optimal guidance \\n discount factor', color=\"tab:green\")\n",
        "\n",
        "\n",
        "  # place a text box in upper left in axes coords\n",
        "  props = dict(boxstyle='round', facecolor='white', alpha=0.5)\n",
        "  textstr = ( r'$\\gamma_{\\mathrm{eval}}=%.2f$ ' % (0.99))\n",
        "  ax1.text(0.75, 0.15, textstr, transform=ax1.transAxes, \n",
        "           fontsize=18, fontweight='bold',\n",
        "           verticalalignment='top', bbox=props)\n",
        "\n",
        "  ax1.annotate('Planning Loss ',\n",
        "            xy=(11, 0.82), xycoords='data',\n",
        "            xytext=(14, 0.55), textcoords='data',\n",
        "            arrowprops=dict(facecolor='black', shrink=0.05),\n",
        "            horizontalalignment='right', verticalalignment='center')\n",
        "  \n",
        "  ax2.annotate('Optimal Gamma ($\\gamma^\\star$)',\n",
        "            xy=(10, 0.82), xycoords='data',\n",
        "            xytext=(14, 0.95), textcoords='data',\n",
        "            arrowprops=dict(facecolor='green', shrink=0.05),\n",
        "            horizontalalignment='right', verticalalignment='center')\n",
        "\n",
        "  # plt.savefig((os.path.join(DIR,\"seeds_randomMDP_lossanddiscount_fnoftasks_1STD\"+ \".pdf\")), \n",
        "  #             dpi=200, orientation='portrait', \n",
        "  #             papertype=None, \n",
        "  #             format=None,\n",
        "  #             transparent=True, \n",
        "  #             bbox_inches='tight', \n",
        "  #             pad_inches=0,\n",
        "  #             figsize=(4, 4))\n",
        "  plt.show()\n",
        "  return fig\n",
        "\n",
        "\n",
        "# plot chosen discount in each round t    \n",
        "figomrl_discount = plotdiscount_paper(avg_loss_per_round, \n",
        "                                      std_dev_per_round,\n",
        "                                      gammas, num_tasks)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "MW5P_agYgYuZ"
      },
      "source": [
        "## Adaptation of Planning Horizon - Main script"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "a7abPg2U72-r",
        "cellView": "form"
      },
      "outputs": [],
      "source": [
        "# @title Load Saved Data\n",
        "\n",
        "num_tasks = 15 # @param{\"type\": \"string\"}\n",
        "runs = 100 # @param{\"type\": \"string\"}\n",
        "\n",
        "filename = 'seeds_ours_2022_sacounts_exp_dirich_0PT015_1PT5N_100runs_15tasks_5samples_5trajle.npy' # @param{\"type\": \"string\"}\n",
        "DIR ='/content/drive/My Drive/onlinemetarl_experiments'  # @param{\"type\": \"string\"}\n",
        "\n",
        "gammas = [gamma_ for gamma_ in np.arange(0, 0.99, 0.1)]\n",
        "gammas.append(0.99)\n",
        "trueGamma = max(gammas)\n",
        "\n",
        "\n",
        "fullpath = os.path.join(DIR, filename)\n",
        "with open(fullpath, 'rb') as f:\n",
        "  proto = storage_pb2.Data()\n",
        "  proto.ParseFromString(f.read())\n",
        "  decoded_A = codec.decode(proto)\n",
        "\n",
        "loss_per_run =  np.zeros((runs, num_tasks, len(gammas)))\n",
        "err_per_run  =  np.zeros((runs, num_tasks, len(gammas)))\n",
        "optGamma_run = np.zeros((runs, num_tasks))\n",
        "\n",
        "avg_loss_per_round = np.zeros((num_tasks, len(gammas)))\n",
        "avg_err_per_round = np.zeros((num_tasks, len(gammas)))\n",
        "\n",
        "for i in range(runs):\n",
        "  loss_per_run[i, :, :] = decoded_A[i][0]\n",
        "  err_per_run[i, :, :] = decoded_A[i][1]\n",
        "  optGamma_run[i,:] = decoded_A[i][2]\n",
        "\n",
        "avg_loss_per_round = np.mean(loss_per_run, axis=0)\n",
        "avg_err_per_round = np.mean(err_per_run, axis=0)\n",
        "std_dev_per_round = np.std(loss_per_run, axis=0)\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "lkNFY71K0Vt7",
        "outputId": "b1d36862-6adb-4219-e625-8b87b937b7a8",
        "cellView": "form"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "0.6500936134803411\n",
            "0.7443411464342337\n",
            "0.7961107336956521\n",
            "0.8347071221385196\n",
            "0.8668515950069346\n",
            "0.8947558360449406\n",
            "0.9192335168298524\n",
            "0.9405545980187217\n",
            "0.9587523436168399\n",
            "0.9737572423194897\n",
            "0.985465282953959\n",
            "0.9937772978201496\n",
            "0.9986234360020648\n",
            "0.9986304496690254\n"
          ]
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "<Figure size 640x480 with 1 Axes>"
            ],
            "image/png": "[encoded data, 72660 characters]\n"
          },
          "metadata": {}
        }
      ],
      "source": [
        "# @title Latest Plan Loss: Choice of $\\gamma$\n",
        "\n",
        "DIR ='/content/drive/My Drive/onlinemetarl_experiments'  # @param{\"type\": \"string\"}\n",
        "schedule = 1/5\n",
        "sn.set_context(\"paper\")\n",
        "sn.set_context(\"paper\", font_scale=2.00)\n",
        "plt.rcParams['savefig.facecolor'] = \"0.8\"\n",
        "plt.rcParams['savefig.edgecolor'] = \"black\"\n",
        "plt.rcParams['pdf.fonttype'] = 42\n",
        "plt.rcParams['ps.fonttype'] = 42\n",
        "\n",
        "m=5\n",
        "colorslist = [\"tab:green\", \"tab:blue\", \"tab:red\"]\n",
        "\n",
        "colors = matplotlib.cm.tab20(range(20))\n",
        "\n",
        "bestinhindsight_loss = np.zeros((num_tasks, 1))\n",
        "bestinhindsight_err = np.zeros((num_tasks, 1))\n",
        "bestinhindsight_gamma = np.zeros((num_tasks, 1))\n",
        "\n",
        "bestvariableinhindsight_loss = np.zeros((num_tasks, 1))\n",
        "bestvariableinhindsight_err = np.zeros((num_tasks, 1))\n",
        "bestvariableinhindsight_gamma = np.zeros((num_tasks, 1))\n",
        "\n",
        "bestasourchedule_loss = np.zeros((num_tasks, 1))\n",
        "bestasourchedule_err = np.zeros((num_tasks, 1))\n",
        "\n",
        "best_loss_analytical = np.zeros((num_tasks, 1))\n",
        "best_err_analytical = np.zeros((num_tasks, 1))\n",
        "\n",
        "gammaeval_loss = np.zeros((num_tasks, 1))\n",
        "gammaeval_err = np.zeros((num_tasks, 1))\n",
        "\n",
        "optGamma_ours_param =  optgamma_ourproposal_alpha(15, m=5,\n",
        "                                                  sigma_true=0.0129,\n",
        "                                                  schedule=1/5)\n",
        "\n",
        "optGamma_Eval = gammas[-1]*np.ones((1, num_tasks))\n",
        "\n",
        "optGamma_ours_analytical=  optgamma_analytical_fn(num_tasks, m, num_states=10,\n",
        "                                          sigma=0.0129,\n",
        "                                          gamma_0=0.35)\n",
        "\n",
        "for t in range(1,num_tasks+1):\n",
        "  # ---------------best in hindsight---------------\n",
        "  # for each gamma, avg planning loss across tasks\n",
        "  # find the gamma that minimizes this avg planning loss across task\n",
        "  avgplanloss_overtasks = np.mean(avg_loss_per_round[:t], axis=0, keepdims=True)\n",
        "  avgplanerr_overtasks = np.mean(avg_err_per_round[:t], axis=0, keepdims=True)\n",
        "  bestinhindsight_idx = np.argmin(avgplanloss_overtasks[0,:])\n",
        "  bestinhindsight_loss[t-1,0] = avgplanloss_overtasks[0,bestinhindsight_idx]\n",
        "  bestinhindsight_err[t-1,0] = avgplanerr_overtasks[0,bestinhindsight_idx]\n",
        "  bestinhindsight_gamma[t-1,0] = gammas[bestinhindsight_idx]\n",
        "\n",
        "\n",
        "  # ---------------Best variable discount in hindsight---------------\n",
        "  # Select the best planning loss across tasks.\n",
        "  avgplanloss_overgammas_idx = np.argmin(avg_loss_per_round[:t], axis=1)\n",
        "  bestvariableinhindsight_loss[t-1,0] = np.mean(np.take_along_axis(avg_loss_per_round[:t], \n",
        "                                                                   np.expand_dims(avgplanloss_overgammas_idx, axis=-1), axis=1), axis=0)\n",
        "  bestvariableinhindsight_err[t-1,0] = np.mean(np.take_along_axis(avg_err_per_round[:t], \n",
        "                                                                   np.expand_dims(avgplanloss_overgammas_idx, axis=-1), axis=1), axis=0)\n",
        "  bestvariableinhindsight_gamma[t-1,0] = gammas[avgplanloss_overgammas_idx[t-1]]\n",
        "\n",
        "\n",
        "\n",
        "  # # ---------------Our schedule of gamma---------------\n",
        "  loss_per_round_ours = np.zeros((t,1))\n",
        "  err_per_round_ours = np.zeros((t,1))\n",
        "\n",
        "  for i, ival in enumerate(optGamma_ours_param[0,:t]):\n",
        "    equivalent_gamma_idx = -1\n",
        "    mindiff = 100\n",
        "    for j, jval in enumerate(gammas):\n",
        "      diff = np.abs(ival-jval) \n",
        "      if diff < mindiff:\n",
        "        mindiff = diff\n",
        "        equivalent_gamma_idx = j\n",
        "        \n",
        "    loss_per_round_ours[i,0] = avg_loss_per_round[i,equivalent_gamma_idx]\n",
        "    err_per_round_ours[i,0] = avg_err_per_round[i,equivalent_gamma_idx]\n",
        "\n",
        "  bestasourchedule_loss[t-1,0] = np.mean(loss_per_round_ours, axis=0, keepdims=True)\n",
        "  bestasourchedule_err[t-1,0] = np.mean(err_per_round_ours, axis=0, keepdims=True)\n",
        "\n",
        "\n",
        "  # # ---------------Gamma Eval Schedule---------------\n",
        "  loss_per_round_eval = np.zeros((t,1))\n",
        "  err_per_round_eval  = np.zeros((t,1))\n",
        "\n",
        "  for i, ival in enumerate(optGamma_Eval[0,:t]):\n",
        "    equivalent_gamma_idx = -1\n",
        "    mindiff = 100\n",
        "    for j, jval in enumerate(gammas):\n",
        "      diff = np.abs(ival-jval) \n",
        "      if diff < mindiff:\n",
        "        mindiff = diff\n",
        "        equivalent_gamma_idx = j\n",
        "        \n",
        "    loss_per_round_eval[i,0] = avg_loss_per_round[i,equivalent_gamma_idx]\n",
        "    err_per_round_eval[i,0] = avg_err_per_round[i,equivalent_gamma_idx]\n",
        "\n",
        "  gammaeval_loss[t-1,0] = np.mean(loss_per_round_eval, axis=0, keepdims=True)\n",
        "  gammaeval_err[t-1,0] = np.mean(err_per_round_eval, axis=0, keepdims=True)\n",
        "\n",
        "\n",
        "\n",
        "  # # ---------------Gamma Analytical Schedule---------------\n",
        "  loss_per_round_analytical = np.zeros((t,1))\n",
        "  err_per_round_analytical  = np.zeros((t,1))\n",
        "\n",
        "  for i, ival in enumerate(optGamma_ours_analytical[0,:t]):\n",
        "    equivalent_gamma_idx = -1\n",
        "    mindiff = 100\n",
        "    for j, jval in enumerate(gammas):\n",
        "      diff = np.abs(ival-jval) \n",
        "      if diff < mindiff:\n",
        "        mindiff = diff\n",
        "        equivalent_gamma_idx = j\n",
        "        \n",
        "    loss_per_round_analytical[i,0] = avg_loss_per_round[i,equivalent_gamma_idx]\n",
        "    err_per_round_analytical[i,0] = avg_err_per_round[i,equivalent_gamma_idx]\n",
        "\n",
        "  best_loss_analytical[t-1,0] = np.mean(loss_per_round_analytical, axis=0, keepdims=True)\n",
        "  best_err_analytical[t-1,0] = np.mean(err_per_round_analytical, axis=0, keepdims=True)\n",
        "\n",
        "\n",
        "#--------------------------------------------------------#\n",
        "# plotting the planning loss baselines \n",
        "fig, ax1 = plt.subplots()\n",
        "xtoplot = list(range(1, num_tasks+1))\n",
        "ax1.set_yscale('log')\n",
        "plt.ylim((0.48, 2.9))   # set the ylim to bottom, top\n",
        "# plt.ylim((0.47, 1.1))   # set the ylim to bottom, top #zoomed \n",
        "\n",
        "plt.errorbar(xtoplot,\n",
        "             bestinhindsight_loss,\n",
        "             bestinhindsight_err[:,0], \n",
        "             label='Best Fixed',\n",
        "             color = \"tab:orange\",\n",
        "             linewidth=2.5)\n",
        "\n",
        "plt.errorbar(xtoplot,\n",
        "             bestvariableinhindsight_loss,\n",
        "             bestvariableinhindsight_err[:,0], \n",
        "             label='Dynamic Best',\n",
        "             color = \"tab:blue\",\n",
        "             linewidth=2.5)\n",
        "\n",
        "plt.errorbar(xtoplot,\n",
        "             gammaeval_loss,\n",
        "             gammaeval_err[:,0], \n",
        "             label='$\\gamma_{eval}=0.99$',\n",
        "             color = \"tab:red\",\n",
        "             linewidth=2.5,)\n",
        "\n",
        "plt.errorbar(xtoplot,\n",
        "             bestasourchedule_loss,\n",
        "             bestasourchedule_err[:,0], \n",
        "             label=r'$\\gamma=f({m,\\alpha_{t}, \\sigma_{t}, T})$ ',\n",
        "             color = \"black\",\n",
        "             linewidth=2.5)\n",
        "\n",
        "plt.errorbar(xtoplot,\n",
        "             best_loss_analytical,\n",
        "             best_err_analytical[:,0], \n",
        "             label= r'$\\gamma=\\min\\{1,\\gamma_0+\\tilde\\gamma\\}$ ',\n",
        "             color = 'tab:pink',\n",
        "             linewidth=2.5)\n",
        "\n",
        "\n",
        "ax1.set_xlabel('Number of tasks (T)')\n",
        "ax1.set_ylabel('Planning Loss')\n",
        "ax1.set_xticks(xtoplot)\n",
        "\n",
        "plt.legend(bbox_to_anchor=(1.11, 1.1), fancybox=True, shadow=True, ncol=2,\n",
        "          facecolor='w', fontsize=18)\n",
        "          \n",
        "matplotlib.rc('axes', edgecolor='black')\n",
        "\n",
        "# plt.savefig((os.path.join(DIR,\"lossforgammaschedules_600runs_usingstdbysqrtruns\"+ \".pdf\")), \n",
        "#             dpi=200, orientation='portrait', \n",
        "#             papertype=None, \n",
        "#             format=None,\n",
        "#             transparent=True, \n",
        "#             bbox_inches='tight', \n",
        "#             pad_inches=0,\n",
        "#             figsize=(4, 4))\n",
        "plt.show()\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "UCflKavSzenp",
        "outputId": "64af84f8-950f-4e13-a02a-a6c44ed10598",
        "cellView": "form"
      },
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "<Figure size 640x480 with 1 Axes>"
            ],
            "image/png": "[encoded data, 70900 characters]\n"
          },
          "metadata": {}
        }
      ],
      "source": [
        "# @title Ours for different $\\gamma_0$\n",
        "\n",
        "DIR ='/content/drive/My Drive/onlinemetarl_experiments'  # @param{\"type\": \"string\"}\n",
        "schedule = 1/5\n",
        "sn.set_context(\"paper\")\n",
        "sn.set_context(\"paper\", font_scale=2.00)\n",
        "plt.rcParams['savefig.facecolor'] = \"0.8\"\n",
        "plt.rcParams['savefig.edgecolor'] = \"black\"\n",
        "plt.rcParams['pdf.fonttype'] = 42\n",
        "plt.rcParams['ps.fonttype'] = 42\n",
        "\n",
        "\n",
        "m=5\n",
        "colorslist = [\"tab:green\", \"tab:blue\", \"tab:red\"]\n",
        "\n",
        "\n",
        "\n",
        "gammaeval_loss = np.zeros((num_tasks, 1))\n",
        "gammaeval_err = np.zeros((num_tasks, 1))\n",
        "\n",
        "\n",
        "optGamma_Eval = gammas[-1]*np.ones((1, num_tasks))\n",
        "\n",
        "optGamma_ours_analytical_1 =  optgamma_analytical_fn(num_tasks, m, num_states=10,\n",
        "                                          sigma=0.0129,\n",
        "                                          gamma_0=0.35)\n",
        "best_loss_analytical_1 = np.zeros((num_tasks, 1))\n",
        "best_err_analytical_1 = np.zeros((num_tasks, 1))\n",
        "\n",
        "\n",
        "\n",
        "optGamma_ours_analytical_2 =  optgamma_analytical_fn(num_tasks, m, num_states=10,\n",
        "                                          sigma=0.0129,\n",
        "                                          gamma_0=0.25)\n",
        "best_loss_analytical_2 = np.zeros((num_tasks, 1))\n",
        "best_err_analytical_2 = np.zeros((num_tasks, 1))\n",
        "\n",
        "\n",
        "optGamma_ours_analytical_3 =  optgamma_analytical_fn(num_tasks, m, num_states=10,\n",
        "                                          sigma=0.0129,\n",
        "                                          gamma_0=0.45)\n",
        "best_loss_analytical_3 = np.zeros((num_tasks, 1))\n",
        "best_err_analytical_3 = np.zeros((num_tasks, 1))\n",
        "\n",
        "\n",
        "optGamma_ours_analytical_4 =  optgamma_analytical_fn(num_tasks, m, num_states=10,\n",
        "                                          sigma=0.0129,\n",
        "                                          gamma_0=0.50)\n",
        "best_loss_analytical_4 = np.zeros((num_tasks, 1))\n",
        "best_err_analytical_4 = np.zeros((num_tasks, 1))\n",
        "\n",
        "\n",
        "for t in range(1,num_tasks+1):\n",
        "  # # ---------------Gamma Eval Schedule---------------\n",
        "  loss_per_round_eval = np.zeros((t,1))\n",
        "  err_per_round_eval  = np.zeros((t,1))\n",
        "\n",
        "  for i, ival in enumerate(optGamma_Eval[0,:t]):\n",
        "    equivalent_gamma_idx = -1\n",
        "    mindiff = 100\n",
        "    for j, jval in enumerate(gammas):\n",
        "      diff = np.abs(ival-jval) \n",
        "      if diff < mindiff:\n",
        "        mindiff = diff\n",
        "        equivalent_gamma_idx = j\n",
        "        \n",
        "    loss_per_round_eval[i,0] = avg_loss_per_round[i,equivalent_gamma_idx]\n",
        "    err_per_round_eval[i,0] = std_dev_per_round[i,equivalent_gamma_idx]\n",
        "\n",
        "  gammaeval_loss[t-1,0] = np.mean(loss_per_round_eval, axis=0, keepdims=True)\n",
        "  gammaeval_err[t-1,0] = np.mean(err_per_round_eval, axis=0, keepdims=True)\n",
        "\n",
        "\n",
        "\n",
        "  # # ---------------Gamma Analytical Schedule---------------\n",
        "  loss_per_round_analytical_1 = np.zeros((t,1))\n",
        "  err_per_round_analytical_1   = np.zeros((t,1))\n",
        "\n",
        "  for i, ival in enumerate(optGamma_ours_analytical_1[0,:t]):\n",
        "    equivalent_gamma_idx = -1\n",
        "    mindiff = 100\n",
        "    for j, jval in enumerate(gammas):\n",
        "      diff = np.abs(ival-jval) \n",
        "      if diff < mindiff:\n",
        "        mindiff = diff\n",
        "        equivalent_gamma_idx = j\n",
        "        \n",
        "    loss_per_round_analytical_1 [i,0] = avg_loss_per_round[i,equivalent_gamma_idx]\n",
        "    err_per_round_analytical_1 [i,0] = std_dev_per_round[i,equivalent_gamma_idx]\n",
        "\n",
        "  best_loss_analytical_1 [t-1,0] = np.mean(loss_per_round_analytical_1 , axis=0, keepdims=True)\n",
        "  best_err_analytical_1 [t-1,0] = np.mean(err_per_round_analytical_1 , axis=0, keepdims=True)\n",
        "\n",
        "\n",
        " # # ---------------Gamma Analytical Schedule---------------\n",
        "  loss_per_round_analytical_2 = np.zeros((t,1))\n",
        "  err_per_round_analytical_2   = np.zeros((t,1))\n",
        "\n",
        "  for i, ival in enumerate(optGamma_ours_analytical_2[0,:t]):\n",
        "    equivalent_gamma_idx = -1\n",
        "    mindiff = 100\n",
        "    for j, jval in enumerate(gammas):\n",
        "      diff = np.abs(ival-jval) \n",
        "      if diff < mindiff:\n",
        "        mindiff = diff\n",
        "        equivalent_gamma_idx = j\n",
        "        \n",
        "    loss_per_round_analytical_2 [i,0] = avg_loss_per_round[i,equivalent_gamma_idx]\n",
        "    err_per_round_analytical_2 [i,0] = std_dev_per_round[i,equivalent_gamma_idx]\n",
        "\n",
        "  best_loss_analytical_2 [t-1,0] = np.mean(loss_per_round_analytical_2 , axis=0, keepdims=True)\n",
        "  best_err_analytical_2 [t-1,0] = np.mean(err_per_round_analytical_2 , axis=0, keepdims=True)\n",
        "\n",
        "\n",
        "  # # ---------------Gamma Analytical Schedule---------------\n",
        "  loss_per_round_analytical_3 = np.zeros((t,1))\n",
        "  err_per_round_analytical_3  = np.zeros((t,1))\n",
        "\n",
        "  for i, ival in enumerate(optGamma_ours_analytical_3[0,:t]):\n",
        "    equivalent_gamma_idx = -1\n",
        "    mindiff = 100\n",
        "    for j, jval in enumerate(gammas):\n",
        "      diff = np.abs(ival-jval) \n",
        "      if diff < mindiff:\n",
        "        mindiff = diff\n",
        "        equivalent_gamma_idx = j\n",
        "        \n",
        "    loss_per_round_analytical_3 [i,0] = avg_loss_per_round[i,equivalent_gamma_idx]\n",
        "    err_per_round_analytical_3 [i,0] = std_dev_per_round[i,equivalent_gamma_idx]\n",
        "\n",
        "  best_loss_analytical_3 [t-1,0] = np.mean(loss_per_round_analytical_3 , axis=0, keepdims=True)\n",
        "  best_err_analytical_3 [t-1,0] = np.mean(err_per_round_analytical_3 , axis=0, keepdims=True)\n",
        "\n",
        "\n",
        "\n",
        "  # # ---------------Gamma Analytical Schedule---------------\n",
        "  loss_per_round_analytical_4 = np.zeros((t,1))\n",
        "  err_per_round_analytical_4  = np.zeros((t,1))\n",
        "\n",
        "  for i, ival in enumerate(optGamma_ours_analytical_4[0,:t]):\n",
        "    equivalent_gamma_idx = -1\n",
        "    mindiff = 100\n",
        "    for j, jval in enumerate(gammas):\n",
        "      diff = np.abs(ival-jval) \n",
        "      if diff < mindiff:\n",
        "        mindiff = diff\n",
        "        equivalent_gamma_idx = j\n",
        "        \n",
        "    loss_per_round_analytical_4[i,0] = avg_loss_per_round[i,equivalent_gamma_idx]\n",
        "    err_per_round_analytical_4[i,0] = std_dev_per_round[i,equivalent_gamma_idx]\n",
        "\n",
        "  best_loss_analytical_4 [t-1,0] = np.mean(loss_per_round_analytical_4 , axis=0, keepdims=True)\n",
        "  best_err_analytical_4 [t-1,0] = np.mean(err_per_round_analytical_4 , axis=0, keepdims=True)\n",
        "\n",
        "\n",
        "#--------------------------------------------------------#\n",
        "# plotting the planning loss baselines \n",
        "fig, ax1 = plt.subplots()\n",
        "xtoplot = list(range(1, num_tasks+1))\n",
        "ax1.set_yscale('log')\n",
        "plt.ylim((0.5, 1))   # set the ylim to bottom, top # use this for analytical\n",
        "\n",
        "plt.errorbar(xtoplot,\n",
        "             best_loss_analytical_1,\n",
        "             #(1.96*best_err_analytical_1[:,0])/np.sqrt(runs),\n",
        "             (best_err_analytical_1[:,0])/np.sqrt(runs), \n",
        "             label='$\\gamma_0=0.35$',\n",
        "             color = 'tab:purple',\n",
        "             linewidth=2.5)\n",
        "\n",
        "plt.errorbar(xtoplot,\n",
        "             best_loss_analytical_2,\n",
        "             #(1.96*best_err_analytical_2[:,0])/np.sqrt(runs), \n",
        "             (best_err_analytical_2[:,0])/np.sqrt(runs), \n",
        "             label='$\\gamma_0=0.25$',\n",
        "             color = 'tab:green',\n",
        "             linewidth=2.5)\n",
        "\n",
        "\n",
        "plt.errorbar(xtoplot,\n",
        "             best_loss_analytical_3,\n",
        "             #(1.96*best_err_analytical_3[:,0])/np.sqrt(runs), \n",
        "             (best_err_analytical_3[:,0])/np.sqrt(runs), \n",
        "             label='$\\gamma_0=0.45$',\n",
        "             color = 'tab:blue',\n",
        "             linewidth=2.5)\n",
        "\n",
        "plt.errorbar(xtoplot,\n",
        "             best_loss_analytical_4,\n",
        "             #(1.96*best_err_analytical_4[:,0])/np.sqrt(runs), \n",
        "             (best_err_analytical_4[:,0])/np.sqrt(runs), \n",
        "             label='$\\gamma_0=0.50$',\n",
        "             color = 'tab:pink',\n",
        "             linewidth=2.5)\n",
        "\n",
        "# place a text box in upper left in axes coords\n",
        "props = dict(boxstyle='round', facecolor='white', alpha=0.5)\n",
        "textstr = ( r'$\\gamma=\\min\\{1,\\gamma_0+\\tilde\\gamma\\}$') \n",
        "ax1.text(0.6, 0.7, textstr, transform=ax1.transAxes, \n",
        "          fontsize=18, fontweight='bold',\n",
        "          verticalalignment='top', bbox=props)\n",
        "\n",
        "\n",
        "ax1.set_xlabel('Number of tasks (T)')\n",
        "ax1.set_ylabel('Planning Loss')\n",
        "ax1.set_xticks(xtoplot)\n",
        "plt.legend(bbox_to_anchor=(0.99, 0.99), fancybox=True, shadow=True, ncol=2,\n",
        "          facecolor='w', fontsize=18)\n",
        "          \n",
        "matplotlib.rc('axes', edgecolor='black')\n",
        "\n",
        "# plt.savefig((os.path.join(DIR,\"various_gamma0_600runs_usingstdbysqrtruns\"+ \".pdf\")), \n",
        "#             dpi=200, orientation='portrait', \n",
        "#             papertype=None, \n",
        "#             format=None,\n",
        "#             transparent=True, \n",
        "#             bbox_inches='tight', \n",
        "#             pad_inches=0,\n",
        "#             figsize=(4, 4))\n",
        "plt.show()\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "uuOg5mQ5w176",
        "outputId": "cadc34fc-d599-4be2-a4ea-bb40999e2efb"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "0.6500936134803411\n",
            "0.7443411464342337\n",
            "0.7961107336956521\n",
            "0.8347071221385196\n",
            "0.8668515950069346\n",
            "0.8947558360449406\n",
            "0.9192335168298524\n",
            "0.9405545980187217\n",
            "0.9587523436168399\n",
            "0.9737572423194897\n",
            "0.985465282953959\n",
            "0.9937772978201496\n",
            "0.9986234360020648\n",
            "0.9986304496690254\n"
          ]
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "<Figure size 640x480 with 1 Axes>"
            ],
            "image/png": "[encoded data, 80212 characters]\n"
          },
          "metadata": {}
        }
      ],
      "source": [
        "# @title Latest Plan Loss: Choice of $\\gamma$ 1-STD\n",
        "\n",
        "DIR ='/content/drive/My Drive/onlinemetarl_experiments'  # @param{\"type\": \"string\"}\n",
        "schedule = 1/5\n",
        "sn.set_context(\"paper\")\n",
        "sn.set_context(\"paper\", font_scale=2.00)\n",
        "plt.rcParams['savefig.facecolor'] = \"0.8\"\n",
        "plt.rcParams['savefig.edgecolor'] = \"black\"\n",
        "plt.rcParams['pdf.fonttype'] = 42\n",
        "plt.rcParams['ps.fonttype'] = 42\n",
        "\n",
        "m=5\n",
        "colorslist = [\"tab:green\", \"tab:blue\", \"tab:red\"]\n",
        "\n",
        "colors = matplotlib.cm.tab20(range(20))\n",
        "\n",
        "#std_dev_per_round holds tasks*gammas std dev across all runs.\n",
        "\n",
        "bestinhindsight_loss = np.zeros((num_tasks, 1))\n",
        "bestinhindsight_std = np.zeros((num_tasks, 1))\n",
        "\n",
        "bestinhindsight_gamma = np.zeros((num_tasks, 1))\n",
        "\n",
        "bestvariableinhindsight_loss = np.zeros((num_tasks, 1))\n",
        "bestvariableinhindsight_std = np.zeros((num_tasks, 1))\n",
        "bestvariableinhindsight_gamma = np.zeros((num_tasks, 1))\n",
        "\n",
        "bestasperben_loss = np.zeros((num_tasks, 1))\n",
        "bestasperben_std = np.zeros((num_tasks, 1))\n",
        "\n",
        "bestasourchedule_loss = np.zeros((num_tasks, 1))\n",
        "bestasourchedule_std = np.zeros((num_tasks, 1))\n",
        "\n",
        "best_loss_analytical = np.zeros((num_tasks, 1))\n",
        "best_std_analytical = np.zeros((num_tasks, 1))\n",
        "\n",
        "gammaeval_loss = np.zeros((num_tasks, 1))\n",
        "gammaeval_std = np.zeros((num_tasks, 1))\n",
        "\n",
        "\n",
        "optGamma_ours_param =  optgamma_ourproposal_alpha(15, m=5,\n",
        "                                                  sigma_true=0.0129,\n",
        "                                                  schedule=1/5)\n",
        "\n",
        "optGamma_Eval = gammas[-1]*np.ones((1, num_tasks))\n",
        "\n",
        "optGamma_ours_analytical=  optgamma_analytical_fn(num_tasks, m, num_states=10,\n",
        "                                          sigma=0.0129,\n",
        "                                          gamma_0=0.35)\n",
        "\n",
        "for t in range(1,num_tasks+1):\n",
        "  # ---------------best in hindsight---------------\n",
        "  # for each gamma, avg planning loss across tasks\n",
        "  # find the gamma that minimizes this avg planning loss across task\n",
        "  avgplanloss_overtasks = np.mean(avg_loss_per_round[:t], axis=0, keepdims=True)\n",
        "  avgplanerr_overtasks = np.mean(std_dev_per_round[:t], axis=0, keepdims=True)\n",
        "  bestinhindsight_idx = np.argmin(avgplanloss_overtasks[0,:])\n",
        "  bestinhindsight_loss[t-1,0] = avgplanloss_overtasks[0,bestinhindsight_idx]\n",
        "  bestinhindsight_std[t-1,0] = avgplanerr_overtasks[0,bestinhindsight_idx]\n",
        "  bestinhindsight_gamma[t-1,0] = gammas[bestinhindsight_idx]\n",
        "\n",
        "\n",
        "  # ---------------Best variable discount in hindsight---------------\n",
        "  # Select the best planning loss across tasks.\n",
        "  avgplanloss_overgammas_idx = np.argmin(avg_loss_per_round[:t], axis=1)\n",
        "  bestvariableinhindsight_loss[t-1,0] = np.mean(np.take_along_axis(avg_loss_per_round[:t], \n",
        "                                                                   np.expand_dims(avgplanloss_overgammas_idx, axis=-1), axis=1), axis=0)\n",
        "  bestvariableinhindsight_std[t-1,0] = np.mean(np.take_along_axis(std_dev_per_round[:t], \n",
        "                                                                   np.expand_dims(avgplanloss_overgammas_idx, axis=-1), axis=1), axis=0)\n",
        "  bestvariableinhindsight_gamma[t-1,0] = gammas[avgplanloss_overgammas_idx[t-1]]\n",
        "\n",
        "\n",
        "  # # ---------------Our schedule of gamma---------------\n",
        "  loss_per_round_ours = np.zeros((t,1))\n",
        "  err_per_round_ours = np.zeros((t,1))\n",
        "\n",
        "  for i, ival in enumerate(optGamma_ours_param[0,:t]):\n",
        "    equivalent_gamma_idx = -1\n",
        "    mindiff = 100\n",
        "    for j, jval in enumerate(gammas):\n",
        "      diff = np.abs(ival-jval) \n",
        "      if diff < mindiff:\n",
        "        mindiff = diff\n",
        "        equivalent_gamma_idx = j\n",
        "        \n",
        "    loss_per_round_ours[i,0] = avg_loss_per_round[i,equivalent_gamma_idx]\n",
        "    err_per_round_ours[i,0] = std_dev_per_round[i,equivalent_gamma_idx]\n",
        "\n",
        "  bestasourchedule_loss[t-1,0] = np.mean(loss_per_round_ours, axis=0, keepdims=True)\n",
        "  bestasourchedule_std[t-1,0] = np.mean(err_per_round_ours, axis=0, keepdims=True)\n",
        "\n",
        "\n",
        "  # # ---------------Gamma Eval Schedule---------------\n",
        "  loss_per_round_eval = np.zeros((t,1))\n",
        "  err_per_round_eval  = np.zeros((t,1))\n",
        "\n",
        "  for i, ival in enumerate(optGamma_Eval[0,:t]):\n",
        "    equivalent_gamma_idx = -1\n",
        "    mindiff = 100\n",
        "    for j, jval in enumerate(gammas):\n",
        "      diff = np.abs(ival-jval) \n",
        "      if diff < mindiff:\n",
        "        mindiff = diff\n",
        "        equivalent_gamma_idx = j\n",
        "        \n",
        "    loss_per_round_eval[i,0] = avg_loss_per_round[i,equivalent_gamma_idx]\n",
        "    err_per_round_eval[i,0] = std_dev_per_round[i,equivalent_gamma_idx]\n",
        "\n",
        "  gammaeval_loss[t-1,0] = np.mean(loss_per_round_eval, axis=0, keepdims=True)\n",
        "  gammaeval_std[t-1,0] = np.mean(err_per_round_eval, axis=0, keepdims=True)\n",
        "\n",
        "\n",
        "\n",
        "  # # ---------------Gamma Analytical Schedule---------------\n",
        "  loss_per_round_analytical = np.zeros((t,1))\n",
        "  err_per_round_analytical  = np.zeros((t,1))\n",
        "\n",
        "  for i, ival in enumerate(optGamma_ours_analytical[0,:t]):\n",
        "    equivalent_gamma_idx = -1\n",
        "    mindiff = 100\n",
        "    for j, jval in enumerate(gammas):\n",
        "      diff = np.abs(ival-jval) \n",
        "      if diff < mindiff:\n",
        "        mindiff = diff\n",
        "        equivalent_gamma_idx = j\n",
        "        \n",
        "    loss_per_round_analytical[i,0] = avg_loss_per_round[i,equivalent_gamma_idx]\n",
        "    err_per_round_analytical[i,0] = std_dev_per_round[i,equivalent_gamma_idx]\n",
        "\n",
        "  best_loss_analytical[t-1,0] = np.mean(loss_per_round_analytical, axis=0, keepdims=True)\n",
        "  best_std_analytical[t-1,0] = np.mean(err_per_round_analytical, axis=0, keepdims=True)\n",
        "\n",
        "\n",
        "#--------------------------------------------------------#\n",
        "# plotting the planning loss baselines \n",
        "fig, ax1 = plt.subplots()\n",
        "xtoplot = list(range(1, num_tasks+1))\n",
        "ax1.set_yscale('log')\n",
        "#plt.ylim((0.52, 2.9))\n",
        "plt.ylim((0.5, 1.0))   # set the ylim to bottom, top #zoomed \n",
        "\n",
        "plt.errorbar(xtoplot,\n",
        "             bestinhindsight_loss,\n",
        "             #(1.96*bestinhindsight_std[:,0])/np.sqrt(runs),\n",
        "             (bestinhindsight_std[:,0])/np.sqrt(runs), \n",
        "             label='Best Fixed',\n",
        "             color = \"tab:orange\",\n",
        "             linewidth=2.5)\n",
        "\n",
        "plt.errorbar(xtoplot,\n",
        "             bestvariableinhindsight_loss,\n",
        "             #(1.96*bestvariableinhindsight_std[:,0])/np.sqrt(runs),\n",
        "             (bestvariableinhindsight_std[:,0])/np.sqrt(runs), \n",
        "             label='Dynamic Best',\n",
        "             color = \"tab:blue\",\n",
        "             linewidth=2.5)\n",
        "\n",
        "plt.errorbar(xtoplot,\n",
        "             bestasourchedule_loss,\n",
        "             #(1.96*bestasourchedule_std[:,0])/np.sqrt(runs), \n",
        "             (bestasourchedule_std[:,0])/np.sqrt(runs), \n",
        "             label=r'$\\gamma=f({m,\\alpha_{t}, \\sigma_{t}, T})$ ',\n",
        "             color = \"black\",\n",
        "             linewidth=2.5)\n",
        "\n",
        "plt.errorbar(xtoplot,\n",
        "             best_loss_analytical,\n",
        "             #(1.96*best_std_analytical[:,0])/np.sqrt(runs),\n",
        "             (best_std_analytical[:,0])/np.sqrt(runs), \n",
        "             label= r'$\\gamma=\\min\\{1,\\gamma_0+\\tilde\\gamma\\}$ ',\n",
        "             color = 'tab:pink',\n",
        "             linewidth=2.5)\n",
        "\n",
        "\n",
        "ax1.set_xlabel('Number of tasks (T)')\n",
        "ax1.set_ylabel('Planning Loss')\n",
        "ax1.set_xticks(xtoplot)\n",
        "plt.legend(bbox_to_anchor=(1.11, 1.1), fancybox=True, shadow=True, ncol=2,\n",
        "          facecolor='w', fontsize=18)\n",
        "          \n",
        "matplotlib.rc('axes', edgecolor='black')\n",
        "\n",
        "# plt.savefig((os.path.join(DIR,\"lossforgammaschedules_600runs_zoomed_usingstdbysqrtruns\"+ \".pdf\")), \n",
        "#             dpi=200, orientation='portrait', \n",
        "#             papertype=None, \n",
        "#             format=None,\n",
        "#             transparent=True, \n",
        "#             bbox_inches='tight', \n",
        "#             pad_inches=0,\n",
        "#             figsize=(4, 4))\n",
        "# plt.show()\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "KspSIO6c4MTM"
      },
      "source": [
        "# Run Different Baselines - Save and Compile Results\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "aCnEYwQu9MlT",
        "cellView": "form"
      },
      "outputs": [],
      "source": [
        "# @title  Baseline Methods.\n",
        "def compute_optimal_alpha(sigma0_1, sigma0_2, samples_):\n",
        "  '''\n",
        "  alpha* =  \\frac{1}{\\sigma^2_o * m + 1}\n",
        "  '''\n",
        "\n",
        "  alpha1 = 0.01*np.ones((sigma0_1.shape[0], 1))\n",
        "  alpha2 = 0.01*np.ones((sigma0_2.shape[0], 1))\n",
        "\n",
        "  for i in range(sigma0_1.shape[0]):\n",
        "    sigma0_1_square = np.square(np.max(sigma0_1[i, :]))\n",
        "    sigma0_2_square = np.square(np.max(sigma0_2[i, :]))\n",
        "    alpha1[i,0] = 1/(np.multiply(sigma0_1_square, samples_) + 1)   \n",
        "    alpha2[i,0] = 1/(np.multiply(sigma0_2_square, samples_) + 1)   \n",
        "\n",
        "  assert np.all((alpha1 >=0) & (alpha1 <=1) )\n",
        "  assert np.all((alpha2 >=0) & (alpha2 <=1) )\n",
        "\n",
        "  return alpha1, alpha2\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "def update_alpha_knownsigma(sigma0_1, sigma0_2, samples_,\n",
        "                            v2, roundt):\n",
        "  '''\n",
        "  alpha_t =  \\frac{v^2}{\\sigma^2_o * (1+1/t) * m + v^2}\n",
        "  '''\n",
        "\n",
        "  alpha1_t = 0.01*np.ones((sigma0_1.shape[0], 1))\n",
        "  alpha2_t = 0.01*np.ones((sigma0_2.shape[0], 1))\n",
        "\n",
        "  for i in range(sigma0_1.shape[0]):\n",
        "    alpha1_t[i,0] = v2/(np.multiply(np.multiply(np.max(sigma0_1[i, :]), (1+1/roundt)), samples_) + v2)   #make v2 smaller, and add (1-1/np.sqrt(roundt)) as a correction term multiplier.\n",
        "    alpha2_t[i,0] = v2/(np.multiply(np.multiply(np.max(sigma0_2[i, :]), (1+1/roundt)), samples_) + v2) \n",
        "\n",
        "  assert np.all((alpha1_t >=0) & (alpha1_t <=1) )\n",
        "  assert np.all((alpha2_t >=0) & (alpha2_t <=1) )\n",
        "\n",
        "  return alpha1_t, alpha2_t\n",
        "\n",
        "\n",
        "def experiment_omrl_lowerboundknownsigma(num_tasks, gammas, num_states, k, samples_, trials,\n",
        "                               trueGamma, v2, max_trajectory_length,\n",
        "                               seed_in):\n",
        "  # planning loss per round t\n",
        "  loss_per_round =  np.zeros((num_tasks, len(gammas)))\n",
        "  err_per_round  =  np.zeros((num_tasks, len(gammas)))\n",
        "  optGamma_round = np.zeros((1, num_tasks))\n",
        "  var0_gap_per_round = np.zeros((1, num_tasks))\n",
        "  var1_gap_per_round = np.zeros((1, num_tasks))\n",
        "  mean0_gap_per_round = np.zeros((1, num_tasks))\n",
        "  mean1_gap_per_round = np.zeros((1, num_tasks))\n",
        "\n",
        "\n",
        "  # initialize alpha\n",
        "  alpha1_t = 0.0 * np.ones((num_states,1))\n",
        "  alpha2_t = 0.0 * np.ones((num_states,1))\n",
        "  alpha_t = [alpha1_t, alpha2_t]\n",
        "\n",
        "  #Run specific seed\n",
        "  np.random.seed(seed_in)\n",
        "\n",
        "  # Generate prior distribution from randomMDP\n",
        "  P1_Prior, R1 = randomMDP(num_states, k)\n",
        "  P2_Prior, R2 = randomMDP(num_states, k)\n",
        "\n",
        "  sigma0_1 = np.zeros((num_states, num_states))\n",
        "  sigma0_2 = np.zeros((num_states, num_states))\n",
        "\n",
        "  # define the dirichlet prior distribution\n",
        "  # todo: try different, 1.5/2 and 1.5/4 and the expectation is u will not decrease too fast\n",
        "  # more or less look similar. # alpha* will not increase too much. \n",
        "  alphasdirich1 = (P1_Prior + 0.015)*1.5*num_states\n",
        "  alphasdirich2 = (P2_Prior + 0.015)*1.5*num_states\n",
        "\n",
        "  # ground truth variance of prior distribution\n",
        "  for i in range(0, num_states):\n",
        "    z = np.sum(alphasdirich1[i,:])\n",
        "    sigma0_1[i,:] = np.multiply(alphasdirich1[i,:]/z, (1-alphasdirich1[i,:]/z))\n",
        "    sigma0_1[i,:] = sigma0_1[i,:]/(z+1)\n",
        "\n",
        "    z = np.sum(alphasdirich2[i,:])\n",
        "    sigma0_2[i,:] = np.multiply(alphasdirich2[i,:]/z, (1-alphasdirich2[i,:]/z))\n",
        "    sigma0_2[i,:] = sigma0_2[i,:]/(z+1)\n",
        "\n",
        "\n",
        "  # set meta-initialize P_o_hat, sigma_o_hat, alpha_t > 0\n",
        "  # the learner knows the true mean Po - centre of task ball.\n",
        "  P1ohat_t = P1_Prior\n",
        "  P2ohat_t = P2_Prior\n",
        "\n",
        "  # the learner knows the true variance of the task ball.\n",
        "  sigmao1hat_t = sigma0_1\n",
        "  sigmao2hat_t = sigma0_2\n",
        "\n",
        "  for roundt in range(1, num_tasks+1):\n",
        "      # sample a Pt ~ task dist. for round t\n",
        "      Pt_1 = np.zeros((num_states, num_states))\n",
        "      Pt_2 = np.zeros((num_states, num_states))  \n",
        "      for i in range(0, num_states):\n",
        "        Pt_1[i,:] = np.random.dirichlet(alphasdirich1[i,:])\n",
        "        Pt_2[i,:] = np.random.dirichlet(alphasdirich2[i,:])\n",
        "      _unit_test_P(Pt_1)\n",
        "      _unit_test_P(Pt_2) \n",
        "      \n",
        "      # generate a batch of data and process to compute planning loss\n",
        "      loss_t, err_t = processwithintask(Pt_1, R1, Pt_2, R2, samples_, \n",
        "                                    gammas, trueGamma, trials,\n",
        "                                    alpha=alpha_t, P1o=P1ohat_t,\n",
        "                                    P2o=P2ohat_t)\n",
        "                                    \n",
        "      idx = np.argmin(loss_t)\n",
        "      optGamma_round[0, roundt-1] = gammas[idx]\n",
        "\n",
        "      loss_per_round[roundt-1, :] = loss_t\n",
        "      err_per_round[roundt-1, :] = err_t\n",
        "\n",
        "      \n",
        "      alpha1_tp1, alpha2_tp1 = update_alpha_knownsigma(sigma0_1, sigma0_2,\n",
        "                                                       samples_, v2,\n",
        "                                                       roundt)\n",
        "      alpha1_t = alpha1_tp1\n",
        "      alpha2_t = alpha2_tp1\n",
        "      alpha_t = [alpha1_t, alpha2_t]    \n",
        "\n",
        "      if roundt==1:\n",
        "        continue\n",
        "\n",
        "      var0_gap_per_round[0, roundt-1] = np.linalg.norm(sigma0_1-sigmao1hat_t, ord=2)\n",
        "      var1_gap_per_round[0, roundt-1] = np.linalg.norm(sigma0_2-sigmao2hat_t, ord=2)\n",
        "      mean0_gap_per_round[0, roundt-1] = np.linalg.norm(P1_Prior-P1ohat_t, ord=2)\n",
        "      mean1_gap_per_round[0, roundt-1] = np.linalg.norm(P2_Prior-P2ohat_t, ord=2)\n",
        "\n",
        "  return loss_per_round, err_per_round, optGamma_round, var0_gap_per_round, var1_gap_per_round, mean0_gap_per_round, mean1_gap_per_round\n",
        "\n",
        "def experiment_omrl_lowerbound(num_tasks, gammas, num_states, k, samples_, trials,\n",
        "                               trueGamma, v2, max_trajectory_length,\n",
        "                               seed_in):\n",
        "  # planning loss per round t\n",
        "  loss_per_round =  np.zeros((num_tasks, len(gammas)))\n",
        "  err_per_round  =  np.zeros((num_tasks, len(gammas)))\n",
        "  optGamma_round = np.zeros((1, num_tasks))\n",
        "  var0_gap_per_round = np.zeros((1, num_tasks))\n",
        "  var1_gap_per_round = np.zeros((1, num_tasks))\n",
        "  mean0_gap_per_round = np.zeros((1, num_tasks))\n",
        "  mean1_gap_per_round = np.zeros((1, num_tasks))\n",
        "\n",
        "\n",
        "  # initialize alpha\n",
        "  alpha1_t = 0.0 * np.ones((num_states,1))\n",
        "  alpha2_t = 0.0 * np.ones((num_states,1))\n",
        "  alpha_t = [alpha1_t, alpha2_t]\n",
        "\n",
        "  #Run specific seed\n",
        "  np.random.seed(seed_in)\n",
        "\n",
        "  # Generate prior distribution from randomMDP\n",
        "  P1_Prior, R1 = randomMDP(num_states, k)\n",
        "  P2_Prior, R2 = randomMDP(num_states, k)\n",
        "\n",
        "  sigma0_1 = np.zeros((num_states, num_states))\n",
        "  sigma0_2 = np.zeros((num_states, num_states))\n",
        "\n",
        "  # define the dirichlet prior distribution\n",
        "  alphasdirich1 = (P1_Prior + 0.015)*1.5*num_states\n",
        "  alphasdirich2 = (P2_Prior + 0.015)*1.5*num_states\n",
        "\n",
        "  # ground truth variance of prior distribution\n",
        "  for i in range(0, num_states):\n",
        "    z = np.sum(alphasdirich1[i,:])\n",
        "    sigma0_1[i,:] = np.multiply(alphasdirich1[i,:]/z, (1-alphasdirich1[i,:]/z))\n",
        "    sigma0_1[i,:] = sigma0_1[i,:]/(z+1)\n",
        "\n",
        "    z = np.sum(alphasdirich2[i,:])\n",
        "    sigma0_2[i,:] = np.multiply(alphasdirich2[i,:]/z, (1-alphasdirich2[i,:]/z))\n",
        "    sigma0_2[i,:] = sigma0_2[i,:]/(z+1)\n",
        "\n",
        "  # set meta-initialize P_o_hat, sigma_o_hat, alpha_t > 0\n",
        "  # the learner knows the true mean Po - centre of task ball.\n",
        "  P1ohat_t = P1_Prior\n",
        "  P2ohat_t = P2_Prior\n",
        "\n",
        "  # the learner knows the true variance of the task ball.\n",
        "  sigmao1hat_t = sigma0_1\n",
        "  sigmao2hat_t = sigma0_2\n",
        "\n",
        "  # the learner knows the optimal alpha parameter.\n",
        "  alpha1_t, alpha2_t = compute_optimal_alpha(sigmao1hat_t, \n",
        "                                             sigmao2hat_t, \n",
        "                                             samples_)\n",
        "  \n",
        "  alpha_t = [alpha1_t, alpha2_t]\n",
        "\n",
        "  for roundt in range(1, num_tasks+1):\n",
        "      # sample a Pt ~ task dist. for round t\n",
        "      Pt_1 = np.zeros((num_states, num_states))\n",
        "      Pt_2 = np.zeros((num_states, num_states))  \n",
        "      for i in range(0, num_states):\n",
        "        Pt_1[i,:] = np.random.dirichlet(alphasdirich1[i,:])\n",
        "        Pt_2[i,:] = np.random.dirichlet(alphasdirich2[i,:])\n",
        "      _unit_test_P(Pt_1)\n",
        "      _unit_test_P(Pt_2) \n",
        "      \n",
        "      # generate a batch of data and process to compute planning loss\n",
        "      loss_t, err_t = processwithintask(Pt_1, R1, Pt_2, R2, samples_, \n",
        "                                    gammas, trueGamma, trials,\n",
        "                                    alpha=alpha_t, P1o=P1ohat_t,\n",
        "                                    P2o=P2ohat_t)\n",
        "                                    \n",
        "      idx = np.argmin(loss_t)\n",
        "      optGamma_round[0, roundt-1] = gammas[idx]\n",
        "\n",
        "      loss_per_round[roundt-1, :] = loss_t\n",
        "      err_per_round[roundt-1, :] = err_t\n",
        "    \n",
        "\n",
        "      if roundt==1:\n",
        "        continue\n",
        "\n",
        "      var0_gap_per_round[0, roundt-1] = np.linalg.norm(sigma0_1-sigmao1hat_t, ord=2)\n",
        "      var1_gap_per_round[0, roundt-1] = np.linalg.norm(sigma0_2-sigmao2hat_t, ord=2)\n",
        "      mean0_gap_per_round[0, roundt-1] = np.linalg.norm(P1_Prior-P1ohat_t, ord=2)\n",
        "      mean1_gap_per_round[0, roundt-1] = np.linalg.norm(P2_Prior-P2ohat_t, ord=2)\n",
        "\n",
        "  return loss_per_round, err_per_round, optGamma_round, var0_gap_per_round, var1_gap_per_round, mean0_gap_per_round, mean1_gap_per_round\n",
        "\n",
        "\n",
        "\n",
        "def experiment_without_metalearning(num_tasks, gammas, num_states, k, samples_, trials,\n",
        "                                    trueGamma, v2, max_trajectory_length,\n",
        "                                    seed_in):\n",
        "  # planning loss per round t\n",
        "  loss_per_round =  np.zeros((num_tasks, len(gammas)))\n",
        "  err_per_round  =  np.zeros((num_tasks, len(gammas)))\n",
        "  optGamma_round = np.zeros((1, num_tasks))\n",
        "\n",
        "  #Run specific seed\n",
        "  np.random.seed(seed_in)\n",
        " \n",
        "  # Generate prior distribution from randomMDP\n",
        "  P1_Prior, R1 = randomMDP(num_states, k)\n",
        "  P2_Prior, R2 = randomMDP(num_states, k)\n",
        "\n",
        "  sigma0_1 = np.zeros((num_states, num_states))\n",
        "  sigma0_2 = np.zeros((num_states, num_states))\n",
        "\n",
        "  # define the dirichlet prior distribution\n",
        "  alphasdirich1 = (P1_Prior + 0.015)*1.5*num_states\n",
        "  alphasdirich2 = (P2_Prior + 0.015)*1.5*num_states\n",
        "\n",
        "  # ground truth variance of prior distribution\n",
        "  for i in range(0, num_states):\n",
        "    z = np.sum(alphasdirich1[i,:])\n",
        "    sigma0_1[i,:] = np.multiply(alphasdirich1[i,:]/z, (1-alphasdirich1[i,:]/z))\n",
        "    sigma0_1[i,:] = sigma0_1[i,:]/(z+1)\n",
        "\n",
        "    z = np.sum(alphasdirich2[i,:])\n",
        "    sigma0_2[i,:] = np.multiply(alphasdirich2[i,:]/z, (1-alphasdirich2[i,:]/z))\n",
        "    sigma0_2[i,:] = sigma0_2[i,:]/(z+1)\n",
        "\n",
        "  for roundt in range(1, num_tasks+1):\n",
        "      # sample a Pt ~ task dist. for round t\n",
        "      Pt_1 = np.zeros((num_states, num_states))\n",
        "      Pt_2 = np.zeros((num_states, num_states))  \n",
        "      for i in range(0, num_states):\n",
        "        Pt_1[i,:] = np.random.dirichlet(alphasdirich1[i,:])\n",
        "        Pt_2[i,:] = np.random.dirichlet(alphasdirich2[i,:])\n",
        "      _unit_test_P(Pt_1)\n",
        "      _unit_test_P(Pt_2) \n",
        "      \n",
        "      # generate a batch of data and process to compute planning loss\n",
        "      loss_t, err_t = processwithintask(Pt_1, R1, Pt_2, R2, samples_,\n",
        "                                        gammas, trueGamma, trials)\n",
        "                                    \n",
        "      idx = np.argmin(loss_t)\n",
        "      optGamma_round[0, roundt-1] = gammas[idx]\n",
        "\n",
        "      loss_per_round[roundt-1, :] = loss_t\n",
        "      err_per_round[roundt-1, :] = err_t\n",
        "\n",
        "  return loss_per_round, err_per_round, optGamma_round"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "PnQdvzy3yH6J",
        "cellView": "form"
      },
      "outputs": [],
      "source": [
        "# @title Launch Runs - Baseline\n",
        "seed = 88\n",
        "\n",
        "samples_ = 5\n",
        "num_tasks = 15\n",
        "num_states = 10\n",
        "max_trajectory_length = 5\n",
        "k = 5\n",
        "v2 = 1/4\n",
        "gammas = [gamma_ for gamma_ in np.arange(0, 0.99, 0.1)]\n",
        "gammas.append(0.99)\n",
        "trueGamma = max(gammas)\n",
        "runs = 50\n",
        "trials = 100\n",
        "\n",
        "\n",
        "iterables = []\n",
        "for run_id in range(runs):\n",
        "  iterables.append([num_tasks, gammas, num_states, k, samples_, \n",
        "                            trials, trueGamma, v2, max_trajectory_length, seed+run_id])\n",
        "\n",
        "\n",
        "\n",
        "%timeit\n",
        "\n",
        "# Baseline I - Known Po, sigmao and alpha*\n",
        "# A = pmap(experiment_omrl_lowerbound, iterables, nproc=8)\n",
        "\n",
        "\n",
        "# Baseline II - Without meta-learning.\n",
        "A = pmap(experiment_without_metalearning, iterables, nproc=8)\n",
        "\n",
        "\n",
        "# Baseline - Known Po, Known sigmao and alpha = f(known sigma, t)\n",
        "# A = pmap(experiment_omrl_lowerboundknownsigma, iterables, nproc=8)\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "G9DFh0StyhDB"
      },
      "outputs": [],
      "source": [
        "# @title  Save Data, Compile Runs\n",
        "\n",
        "\n",
        "filename = '2022_sacounts_exp_dirich_0PT015BY4_1PT5N_100runs_15tasks_5samples_5trajle.npy' # @param{\"type\": \"string\"}\n",
        "DIR ='/content/drive/My Drive/onlinemetarl_experiments/'  # @param{\"type\": \"string\"}\n",
        "fullpath = DIR + filename\n",
        "\n",
        "#----------------START --> Use to load only -----------#\n",
        "\n",
        "num_tasks = 15 # @param{\"type\": \"string\"}\n",
        "runs = 100 # @param{\"type\": \"string\"}\n",
        "\n",
        "\n",
        "gammas = [gamma_ for gamma_ in np.arange(0, 0.99, 0.1)]\n",
        "gammas.append(0.99)\n",
        "trueGamma = max(gammas)\n",
        "\n",
        "\n",
        "fullpath = os.path.join(DIR, filename)\n",
        "with open(fullpath, 'rb') as f:\n",
        "  proto = storage_pb2.Data()\n",
        "  proto.ParseFromString(f.read())\n",
        "  decoded_A = codec.decode(proto)\n",
        "#----------------END ---------------------------------#\n",
        "\n",
        "loss_per_run =  np.zeros((runs, num_tasks, len(gammas)))\n",
        "err_per_run  =  np.zeros((runs, num_tasks, len(gammas)))\n",
        "optGamma_run = np.zeros((runs, num_tasks))\n",
        "\n",
        "avg_loss_per_round = np.zeros((num_tasks, len(gammas)))\n",
        "avg_err_per_round = np.zeros((num_tasks, len(gammas)))\n",
        "std_dev_per_round = np.zeros((num_tasks, len(gammas)))\n",
        "\n",
        "for i in range(runs):\n",
        "  loss_per_run[i, :, :] = decoded_A[i][0]\n",
        "  err_per_run[i, :, :] = decoded_A[i][1]\n",
        "  optGamma_run[i,:] = decoded_A[i][2]\n",
        "\n",
        "avg_loss_per_round = np.mean(loss_per_run, axis=0)\n",
        "avg_err_per_round = np.mean(err_per_run, axis=0)\n",
        "std_dev_per_round = np.std(loss_per_run, axis=0) # standard deviation across runs."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 618
        },
        "id": "tWWur5E5zUMe",
        "outputId": "e14fd29e-5479-4d25-cc51-d97ec8a71cb9"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "0.949638347757141 5\n",
            "1.2045159018301461 7\n",
            "0.9926812590884482 6\n",
            "1.0690266251983802 8\n",
            "0.8325349898365892 7\n",
            "0.8858695988168768 8\n",
            "0.9241748117251446 9\n",
            "0.7966320794113655 8\n"
          ]
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "<Figure size 640x480 with 1 Axes>"
            ],
            "image/png": "[encoded data, 62820 characters]\n"
          },
          "metadata": {}
        }
      ],
      "source": [
        "# @title Paper OMRL Baseline Plots\n",
        "\n",
        "\n",
        "import matplotlib as mpl\n",
        "from matplotlib.ticker import FormatStrFormatter\n",
        "\n",
        "sn.set_context(\"paper\")\n",
        "sn.set_context(\"paper\", font_scale=2.00)\n",
        "plt.rcParams['savefig.facecolor'] = \"0.8\"\n",
        "plt.rcParams['savefig.edgecolor'] = \"black\"\n",
        "plt.rcParams['pdf.fonttype'] = 42\n",
        "plt.rcParams['ps.fonttype'] = 42\n",
        "\n",
        "\n",
        "nlines=8\n",
        "c = np.arange(1, nlines + 1)\n",
        "norm = mpl.colors.Normalize(vmin=c.min()-1, vmax=c.max()+1)\n",
        "cmap = mpl.cm.ScalarMappable(norm=norm, cmap=mpl.cm.bwr)\n",
        "cmap.set_array([])\n",
        "\n",
        "\n",
        "def plotloss_paper(y, e, gammas, numtasks):\n",
        "  fig, ax1 = plt.subplots()\n",
        "  idx = 0\n",
        "  gammalist = []\n",
        "  for i in range(0, numtasks):\n",
        "    if i==0 or i==1 or i==2 or i==4 or i==5 or i==8 or i>=13:  \n",
        "      yk = y[i,:].reshape((1,y.shape[-1]))\n",
        "      ek = e[i,:].reshape((1,y.shape[-1]))\n",
        "      # Use this for 1-STD std_dev_per_round\n",
        "      plt.errorbar(gammas, yk[0,:], yerr=(ek[0,:]/np.sqrt(runs)),\n",
        "                  label = 'T={}'.format(int(i+1)),\n",
        "                  color=cmap.to_rgba(idx + i),\n",
        "                   linewidth=2.5)\n",
        "      min_x = np.argmin(yk[0,:])\n",
        "      min_y = np.min(yk[0,:])\n",
        "      print(min_y, min_x)\n",
        "      plt.scatter(gammas[min_x], min_y, \n",
        "                  color=cmap.to_rgba(idx + i), \n",
        "                  marker='*',\n",
        "                  s=380)\n",
        "      plt.ylim((0, 4))   # set the ylim to bottom, top\n",
        "      plt.xlim((gammas[0], 1))   # set the ylim to bottom, top\n",
        "      idx += 1\n",
        "      gammalist.append(gammas[min_x])\n",
        "  ax1.set_xlabel('$\\gamma$')\n",
        "  ax1.set_ylabel('Planning Loss')\n",
        "  ax1.set_xticks(gammas)\n",
        "  plt.legend(bbox_to_anchor=(1, 1), fancybox=True, shadow=True, ncol=4,\n",
        "          facecolor='w', fontsize=16)\n",
        "          \n",
        "  matplotlib.rc('axes', edgecolor='black')\n",
        "  ax1.xaxis.set_major_formatter(FormatStrFormatter('%.1f'))\n",
        "\n",
        "  # plt.savefig((os.path.join(DIR,\"all_randomMDP_averageplanloss_ours_gamma_markers_bwr_1STD\"+ \".pdf\")), \n",
        "  #             dpi=200, orientation='portrait', \n",
        "  #             papertype=None, \n",
        "  #             format=None,\n",
        "  #             transparent=True, \n",
        "  #             bbox_inches='tight', \n",
        "  #             pad_inches=0,\n",
        "  #             figsize=(4, 4))\n",
        "  plt.show()\n",
        "  return fig, gammalist\n",
        "\n",
        "# planning loss per round t\n",
        "loss_per_run =  np.zeros((runs, num_tasks, len(gammas)))\n",
        "err_per_run  =  np.zeros((runs, num_tasks, len(gammas)))\n",
        "\n",
        "optGamma_run = np.zeros((runs, num_tasks))\n",
        "\n",
        "\n",
        "# track estimates per round t\n",
        "var0_gap_run =  np.zeros((runs, num_tasks))\n",
        "var1_gap_run = np.zeros((runs, num_tasks))\n",
        "mean0_gap_run =  np.zeros((runs, num_tasks))\n",
        "mean1_gap_run =  np.zeros((runs, num_tasks))\n",
        "\n",
        "for i in range(runs):\n",
        "  loss_per_run[i, :, :] = decoded_A[i][0]\n",
        "  err_per_run[i, :, :] = decoded_A[i][1]\n",
        "  optGamma_run[i,:] = decoded_A[i][2] \n",
        "\n",
        "avg_loss_per_round = np.mean(loss_per_run, axis=0)\n",
        "avg_err_per_round = np.mean(err_per_run, axis=0)\n",
        "std_dev_per_round = np.std(loss_per_run, axis=0)\n",
        "\n",
        "\n",
        "# plot planning loss in each round t vs gamma   \n",
        "figomrl_loss, chosen_discount = plotloss_paper(avg_loss_per_round, \n",
        "                                               std_dev_per_round,\n",
        "                                               gammas, num_tasks)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "VD41EXO6eeyY"
      },
      "source": [
        "# Accumulate all Baselines- Plot Averaged Planning Loss\n",
        "\n",
        "Baselines:\n",
        "1.   pomrl\n",
        "2.   ada-pomrl\n",
        "3.   Without Meta Learning\n",
        "4.   Oracle Prior Knowledge\n",
        "5.   Aggregation"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "eR7Adx1JN8fM"
      },
      "outputs": [],
      "source": [
        "# -----------Baeline Specific Variables ------------\n",
        "avg_loss_per_round_dict = {\n",
        "    'ada-POMRL': [],\n",
        "    'Oracle Prior Knowledge': [],\n",
        "    'Without Meta-Learning': [],\n",
        "    'Aggregating': [],\n",
        "    'POMRL': []\n",
        "}\n",
        "avg_err_per_round_dict = {\n",
        "    'ada-POMRL': [],\n",
        "    'Oracle Prior Knowledge': [],\n",
        "    'Without Meta-Learning': [],\n",
        "    'Aggregating': [],\n",
        "    'POMRL': []\n",
        "}"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "7i2b9I2cqJmo"
      },
      "outputs": [],
      "source": [
        "os.chdir(\"/content/drive/My Drive/onlinemetarl_experiments\")\n",
        "!ls"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "PrHYeVMVetW1"
      },
      "outputs": [],
      "source": [
        "# @title  Load Saved Runs, Log Data\n",
        "\n",
        "filename = 'Ours_2022_sacounts_exp_dirich_0PT015BY4_1PT5N_100runs_15tasks_5samples_5trajle.npy' # @param{\"type\": \"string\"}\n",
        "DIR ='/content/drive/My Drive/onlinemetarl_experiments/'  # @param{\"type\": \"string\"}\n",
        "fullpath = DIR + filename\n",
        "\n",
        "#----------------START --> Use to load only -----------#\n",
        "\n",
        "num_tasks = 15 # @param{\"type\": \"string\"}\n",
        "runs = 100 # @param{\"type\": \"string\"}\n",
        "\n",
        "gammas = [gamma_ for gamma_ in np.arange(0, 0.99, 0.1)]\n",
        "gammas.append(0.99)\n",
        "trueGamma = max(gammas)\n",
        "\n",
        "\n",
        "fullpath = os.path.join(DIR, filename)\n",
        "with open(fullpath, 'rb') as f:\n",
        "  proto = storage_pb2.Data()\n",
        "  proto.ParseFromString(f.read())\n",
        "  decoded_A = codec.decode(proto)\n",
        "#----------------END ---------------------------------#\n",
        "\n",
        "\n",
        "# planning loss per round t\n",
        "loss_per_run =  np.zeros((runs, num_tasks, len(gammas)))\n",
        "err_per_run  =  np.zeros((runs, num_tasks, len(gammas)))\n",
        "optGamma_run = np.zeros((runs, num_tasks))\n",
        "\n",
        "for i in range(runs):\n",
        "  loss_per_run[i, :, :] = decoded_A[i][0]\n",
        "  err_per_run[i, :, :] = decoded_A[i][1]\n",
        "  optGamma_run[i,:] = decoded_A[i][2] \n",
        "\n",
        "avg_loss_per_round = np.mean(loss_per_run, axis=0)\n",
        "avg_err_per_round = np.mean(err_per_run, axis=0)\n",
        "std_dev_per_round = np.std(loss_per_run, axis=0)\n",
        "\n",
        "avg_loss_per_round_dict['ada-POMRL'] = avg_loss_per_round\n",
        "avg_err_per_round_dict['ada-POMRL'] = std_dev_per_round\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "CM2A6cpSincq",
        "cellView": "form"
      },
      "outputs": [],
      "source": [
        "# @title Utilites to Plot\n",
        "\n",
        "import matplotlib as mpl\n",
        "from matplotlib.ticker import FormatStrFormatter\n",
        "mpl.rc('axes', edgecolor='black')\n",
        "sn.set_context(\"paper\")\n",
        "sn.set_context(\"paper\", font_scale=1.85)\n",
        "plt.rcParams['savefig.facecolor'] = \"0.8\"\n",
        "plt.rcParams['savefig.edgecolor'] = \"black\"\n",
        "plt.rcParams['pdf.fonttype'] = 42\n",
        "plt.rcParams['ps.fonttype'] = 42\n",
        "\n",
        "\n",
        "\n",
        "nlines=3\n",
        "c = np.arange(1, nlines + 1)\n",
        "norm = mpl.colors.Normalize(vmin=c.min()-1, vmax=c.max()+1)\n",
        "cmap = mpl.cm.ScalarMappable(norm=norm, cmap=mpl.cm.viridis)\n",
        "cmap.set_array([])\n",
        "\n",
        "colorslist = [\"tab:green\", \"tab:blue\", \"tab:red\",  \"tab:purple\", \"tab:orange\"]\n",
        "labels_methods = [r'ada-POMRL', 'Oracle Prior Knowledge', 'Without Meta-Learning', r'POMRL($\\sigma$)'] \n",
        "#labels_methods = [r'ada-POMRL', 'Oracle Prior Knowledge', 'Without Meta-Learning', r'Aggregating($\\alpha=1$)', r'POMRL($\\sigma$)'] \n",
        "\n",
        "def plotloss_fixedgamma(y_dict, e_dict, gammas, \n",
        "                        numtasks, gamma_chosen):\n",
        "  xtoplot = list(range(1, num_tasks+1))\n",
        "  fig, ax1 = plt.subplots()\n",
        "  plt.ylim((0, 4))   # set the ylim to bottom, top\n",
        "  idx_gamma_chosen = gammas.index(gamma_chosen)\n",
        "\n",
        "  method = 0\n",
        "  for key in y_dict:\n",
        "    losslist = []\n",
        "    errlist = []\n",
        "    y = y_dict[key] \n",
        "    e = e_dict[key]\n",
        "    for i in range(0, numtasks):\n",
        "      if i>=0:\n",
        "        yk = y[i,:].reshape((1,y.shape[-1])) #1*11\n",
        "        ek = e[i,:].reshape((1,y.shape[-1])) #1*11\n",
        "        losslist.append(yk[0,idx_gamma_chosen])\n",
        "        errlist.append(ek[0,idx_gamma_chosen])\n",
        "\n",
        "    plt.plot(xtoplot, losslist, color=colorslist[method],\n",
        "            #label = key, \n",
        "            label = labels_methods[method],\n",
        "            linewidth = 2.50,\n",
        "            marker='o', markersize='8')\n",
        "      \n",
        "    losslist_arr = np.array(losslist)\n",
        "    errlist_arr = np.array(errlist)\n",
        "    plt.fill_between(xtoplot, losslist_arr-(errlist_arr/np.sqrt(runs)), losslist_arr+(errlist_arr/np.sqrt(runs)), \n",
        "                    facecolor=colorslist[method], edgecolor=colorslist[method],\n",
        "                    alpha=0.15)\n",
        "    method += 1\n",
        "    \n",
        "\n",
        "  plt.legend(bbox_to_anchor=(1, 1), fancybox=True, shadow=True, ncol=1,\n",
        "          facecolor='w', fontsize=14)\n",
        "\n",
        "  ax1.set_xticks(xtoplot)\n",
        "  ax1.set_xlabel('Task $(T)$')\n",
        "  ax1.set_ylabel('Planning Loss')\n",
        "\n",
        "  # place a text box in upper left in axes coords\n",
        "  props = dict(boxstyle='round', facecolor='white', alpha=0.5)\n",
        "  textstr = ( r'$\\gamma=%.2f= \\gamma_{\\mathrm{eval}}$ ' % (gamma_chosen))\n",
        "  #textstr = ( r'$\\gamma=%.2f$ ' % (gamma_chosen))\n",
        "  ax1.text(0.05, 0.95, textstr, transform=ax1.transAxes, \n",
        "           fontsize=18, fontweight='bold',\n",
        "           verticalalignment='top', bbox=props)\n",
        "\n",
        "  plt.savefig((os.path.join(DIR,\"XYZ_allmethods_fixedgamma_0PT99_1STD_BY4\"+ \".pdf\")), \n",
        "              dpi=200, orientation='portrait', \n",
        "              papertype=None, \n",
        "              format=None,\n",
        "              transparent=True, \n",
        "              bbox_inches='tight', \n",
        "              pad_inches=0,\n",
        "              figsize=(4, 4))\n",
        "  plt.show()\n",
        "  return fig"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 478
        },
        "id": "21tQAcDTfdqn",
        "outputId": "48f3106f-5da6-4f54-fff9-5e9b7508c813"
      },
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "<Figure size 640x480 with 1 Axes>"
            ],
            "image/png": "[encoded data, 93584 characters]\n"
          },
          "metadata": {}
        }
      ],
      "source": [
        "# Planning Loss - Fixed Gamma Value \n",
        "gamma_chosen =  0.99\n",
        "figomrl_loss_fixedgamma = plotloss_fixedgamma(avg_loss_per_round_dict, \n",
        "                                                    avg_err_per_round_dict,\n",
        "                                                    gammas, num_tasks,\n",
        "                                                    gamma_chosen=gamma_chosen)"
      ]
    }
  ],
  "metadata": {
    "accelerator": "TPU",
    "colab": {
      "collapsed_sections": [
        "WfNbd9F5aKN5",
        "r05HA2UWhKz6",
        "AZLAMFTpJKxd",
        "fI_Fo1XHlvW3",
        "OTPraU4oFkIe",
        "3G3q3UouX4HC",
        "DHdEfnmifyt6",
        "OcmMEk_sbFYS",
        "MW5P_agYgYuZ",
        "KspSIO6c4MTM",
        "VD41EXO6eeyY"
      ],
      "machine_shape": "hm",
      "provenance": [],
      "toc_visible": true
    },
    "gpuClass": "standard",
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}