MAML_A2C.ipynb

{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "0u8jyVEHU90_",
        "outputId": "9ebe511b-5dcb-4398-b3cb-a91b1abd8fb5"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Reading package lists... Done\n",
            "Building dependency tree       \n",
            "Reading state information... Done\n",
            "libgl1-mesa-dev is already the newest version (20.0.8-0ubuntu1~18.04.1).\n",
            "libgl1-mesa-dev set to manually installed.\n",
            "software-properties-common is already the newest version (0.96.24.32.18).\n",
            "The following package was automatically installed and is no longer required:\n",
            "  libnvidia-common-460\n",
            "Use 'apt autoremove' to remove it.\n",
            "Suggested packages:\n",
            "  glew-utils\n",
            "The following NEW packages will be installed:\n",
            "  libgl1-mesa-glx libglew-dev libglew2.0 libosmesa6 libosmesa6-dev\n",
            "0 upgraded, 5 newly installed, 0 to remove and 19 not upgraded.\n",
            "Need to get 2,916 kB of archives.\n",
            "After this operation, 12.6 MB of additional disk space will be used.\n",
            "Get:1 http://archive.ubuntu.com/ubuntu bionic-updates/main amd64 libgl1-mesa-glx amd64 20.0.8-0ubuntu1~18.04.1 [5,532 B]\n",
            "Get:2 http://archive.ubuntu.com/ubuntu bionic/universe amd64 libglew2.0 amd64 2.0.0-5 [140 kB]\n",
            "Get:3 http://archive.ubuntu.com/ubuntu bionic/universe amd64 libglew-dev amd64 2.0.0-5 [120 kB]\n",
            "Get:4 http://archive.ubuntu.com/ubuntu bionic-updates/main amd64 libosmesa6 amd64 20.0.8-0ubuntu1~18.04.1 [2,641 kB]\n",
            "Get:5 http://archive.ubuntu.com/ubuntu bionic-updates/main amd64 libosmesa6-dev amd64 20.0.8-0ubuntu1~18.04.1 [8,828 B]\n",
            "Fetched 2,916 kB in 1s (4,071 kB/s)\n",
            "Selecting previously unselected package libgl1-mesa-glx:amd64.\n",
            "(Reading database ... 155680 files and directories currently installed.)\n",
            "Preparing to unpack .../libgl1-mesa-glx_20.0.8-0ubuntu1~18.04.1_amd64.deb ...\n",
            "Unpacking libgl1-mesa-glx:amd64 (20.0.8-0ubuntu1~18.04.1) ...\n",
            "Selecting previously unselected package libglew2.0:amd64.\n",
            "Preparing to unpack .../libglew2.0_2.0.0-5_amd64.deb ...\n",
            "Unpacking libglew2.0:amd64 (2.0.0-5) ...\n",
            "Selecting previously unselected package libglew-dev:amd64.\n",
            "Preparing to unpack .../libglew-dev_2.0.0-5_amd64.deb ...\n",
            "Unpacking libglew-dev:amd64 (2.0.0-5) ...\n",
            "Selecting previously unselected package libosmesa6:amd64.\n",
            "Preparing to unpack .../libosmesa6_20.0.8-0ubuntu1~18.04.1_amd64.deb ...\n",
            "Unpacking libosmesa6:amd64 (20.0.8-0ubuntu1~18.04.1) ...\n",
            "Selecting previously unselected package libosmesa6-dev:amd64.\n",
            "Preparing to unpack .../libosmesa6-dev_20.0.8-0ubuntu1~18.04.1_amd64.deb ...\n",
            "Unpacking libosmesa6-dev:amd64 (20.0.8-0ubuntu1~18.04.1) ...\n",
            "Setting up libosmesa6:amd64 (20.0.8-0ubuntu1~18.04.1) ...\n",
            "Setting up libgl1-mesa-glx:amd64 (20.0.8-0ubuntu1~18.04.1) ...\n",
            "Setting up libglew2.0:amd64 (2.0.0-5) ...\n",
            "Setting up libglew-dev:amd64 (2.0.0-5) ...\n",
            "Setting up libosmesa6-dev:amd64 (20.0.8-0ubuntu1~18.04.1) ...\n",
            "Processing triggers for libc-bin (2.27-3ubuntu1.5) ...\n",
            "Reading package lists... Done\n",
            "Building dependency tree       \n",
            "Reading state information... Done\n",
            "The following package was automatically installed and is no longer required:\n",
            "  libnvidia-common-460\n",
            "Use 'apt autoremove' to remove it.\n",
            "The following NEW packages will be installed:\n",
            "  patchelf\n",
            "0 upgraded, 1 newly installed, 0 to remove and 19 not upgraded.\n",
            "Need to get 46.5 kB of archives.\n",
            "After this operation, 130 kB of additional disk space will be used.\n",
            "Get:1 http://archive.ubuntu.com/ubuntu bionic/universe amd64 patchelf amd64 0.9-1 [46.5 kB]\n",
            "Fetched 46.5 kB in 1s (67.6 kB/s)\n",
            "Selecting previously unselected package patchelf.\n",
            "(Reading database ... 155718 files and directories currently installed.)\n",
            "Preparing to unpack .../patchelf_0.9-1_amd64.deb ...\n",
            "Unpacking patchelf (0.9-1) ...\n",
            "Setting up patchelf (0.9-1) ...\n",
            "Processing triggers for man-db (2.8.3-2ubuntu0.1) ...\n",
            "Looking in indexes: https://pypi.org/simple, https://us-python.pkg.dev/colab-wheels/public/simple/\n",
            "Collecting free-mujoco-py\n",
            "  Downloading free_mujoco_py-2.1.6-py3-none-any.whl (14.1 MB)\n",
            "\u001b[K     |████████████████████████████████| 14.1 MB 5.6 MB/s \n",
            "\u001b[?25hRequirement already satisfied: cffi<2.0.0,>=1.15.0 in /usr/local/lib/python3.7/dist-packages (from free-mujoco-py) (1.15.1)\n",
            "Requirement already satisfied: imageio<3.0.0,>=2.9.0 in /usr/local/lib/python3.7/dist-packages (from free-mujoco-py) (2.9.0)\n",
            "Requirement already satisfied: Cython<0.30.0,>=0.29.24 in /usr/local/lib/python3.7/dist-packages (from free-mujoco-py) (0.29.32)\n",
            "Requirement already satisfied: numpy<2.0.0,>=1.21.3 in /usr/local/lib/python3.7/dist-packages (from free-mujoco-py) (1.21.6)\n",
            "Collecting glfw<2.0.0,>=1.4.0\n",
            "  Downloading glfw-1.12.0-py2.py27.py3.py30.py31.py32.py33.py34.py35.py36.py37.py38-none-manylinux2014_x86_64.whl (203 kB)\n",
            "\u001b[K     |████████████████████████████████| 203 kB 50.3 MB/s \n",
            "\u001b[?25hCollecting fasteners==0.15\n",
            "  Downloading fasteners-0.15-py2.py3-none-any.whl (23 kB)\n",
            "Requirement already satisfied: six in /usr/local/lib/python3.7/dist-packages (from fasteners==0.15->free-mujoco-py) (1.15.0)\n",
            "Collecting monotonic>=0.1\n",
            "  Downloading monotonic-1.6-py2.py3-none-any.whl (8.2 kB)\n",
            "Requirement already satisfied: pycparser in /usr/local/lib/python3.7/dist-packages (from cffi<2.0.0,>=1.15.0->free-mujoco-py) (2.21)\n",
            "Requirement already satisfied: pillow in /usr/local/lib/python3.7/dist-packages (from imageio<3.0.0,>=2.9.0->free-mujoco-py) (7.1.2)\n",
            "Installing collected packages: monotonic, glfw, fasteners, free-mujoco-py\n",
            "Successfully installed fasteners-0.15 free-mujoco-py-2.1.6 glfw-1.12.0 monotonic-1.6\n"
          ]
        }
      ],
      "source": [
        "!apt-get install -y \\\n",
        "    libgl1-mesa-dev \\\n",
        "    libgl1-mesa-glx \\\n",
        "    libglew-dev \\\n",
        "    libosmesa6-dev \\\n",
        "    software-properties-common\n",
        "\n",
        "!apt-get install -y patchelf\n",
        "!pip install free-mujoco-py"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "EvBkQ0-2TIGI"
      },
      "source": [
        "## Cherry TRPO MAML"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "metadata": {
        "id": "wZleUELMTDCf"
      },
      "outputs": [],
      "source": [
        "!pip install cherry-rl learn2learn &> /dev/null"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 3,
      "metadata": {
        "id": "b1cjFZVuS3U2"
      },
      "outputs": [],
      "source": [
        "import random\n",
        "import math\n",
        "import time\n",
        "\n",
        "from copy import deepcopy\n",
        "\n",
        "import cherry as ch\n",
        "import gym\n",
        "import numpy as np\n",
        "import torch\n",
        "from cherry.algorithms import a2c, ppo\n",
        "from cherry.models.robotics import LinearValue\n",
        "from tqdm import tqdm\n",
        "\n",
        "import learn2learn as l2l\n",
        "\n",
        "import torch as th\n",
        "import torch.nn as nn\n",
        "from torch import autograd\n",
        "from torch.distributions.kl import kl_divergence\n",
        "from torch.nn.utils import parameters_to_vector, vector_to_parameters\n",
        "from torch.distributions import Normal, Categorical\n",
        "from torch.utils.tensorboard import SummaryWriter"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "env = gym.make('HalfCheetahForwardBackward-v1')\n",
        "env.reset()"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "U8x8UNw_3uZQ",
        "outputId": "0c5d1919-5ccc-4789-9331-177c4e75e843"
      },
      "execution_count": 4,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Compiling /usr/local/lib/python3.7/dist-packages/mujoco_py/cymj.pyx because it changed.\n",
            "[1/1] Cythonizing /usr/local/lib/python3.7/dist-packages/mujoco_py/cymj.pyx\n",
            "running build_ext\n",
            "building 'mujoco_py.cymj' extension\n",
            "creating /usr/local/lib/python3.7/dist-packages/mujoco_py/generated/_pyxbld_2.0.2.13_37_linuxcpuextensionbuilder\n",
            "creating /usr/local/lib/python3.7/dist-packages/mujoco_py/generated/_pyxbld_2.0.2.13_37_linuxcpuextensionbuilder/temp.linux-x86_64-3.7\n",
            "creating /usr/local/lib/python3.7/dist-packages/mujoco_py/generated/_pyxbld_2.0.2.13_37_linuxcpuextensionbuilder/temp.linux-x86_64-3.7/usr\n",
            "creating /usr/local/lib/python3.7/dist-packages/mujoco_py/generated/_pyxbld_2.0.2.13_37_linuxcpuextensionbuilder/temp.linux-x86_64-3.7/usr/local\n",
            "creating /usr/local/lib/python3.7/dist-packages/mujoco_py/generated/_pyxbld_2.0.2.13_37_linuxcpuextensionbuilder/temp.linux-x86_64-3.7/usr/local/lib\n",
            "creating /usr/local/lib/python3.7/dist-packages/mujoco_py/generated/_pyxbld_2.0.2.13_37_linuxcpuextensionbuilder/temp.linux-x86_64-3.7/usr/local/lib/python3.7\n",
            "creating /usr/local/lib/python3.7/dist-packages/mujoco_py/generated/_pyxbld_2.0.2.13_37_linuxcpuextensionbuilder/temp.linux-x86_64-3.7/usr/local/lib/python3.7/dist-packages\n",
            "creating /usr/local/lib/python3.7/dist-packages/mujoco_py/generated/_pyxbld_2.0.2.13_37_linuxcpuextensionbuilder/temp.linux-x86_64-3.7/usr/local/lib/python3.7/dist-packages/mujoco_py\n",
            "creating /usr/local/lib/python3.7/dist-packages/mujoco_py/generated/_pyxbld_2.0.2.13_37_linuxcpuextensionbuilder/temp.linux-x86_64-3.7/usr/local/lib/python3.7/dist-packages/mujoco_py/gl\n",
            "x86_64-linux-gnu-gcc -pthread -Wno-unused-result -Wsign-compare -DNDEBUG -g -fwrapv -O2 -Wall -g -fstack-protector-strong -Wformat -Werror=format-security -g -fwrapv -O2 -g -fstack-protector-strong -Wformat -Werror=format-security -Wdate-time -D_FORTIFY_SOURCE=2 -fPIC -I/usr/local/lib/python3.7/dist-packages/mujoco_py -I/usr/local/lib/python3.7/dist-packages/mujoco_py/binaries/linux/mujoco210/include -I/usr/local/lib/python3.7/dist-packages/numpy/core/include -I/usr/include/python3.7m -c /usr/local/lib/python3.7/dist-packages/mujoco_py/cymj.c -o /usr/local/lib/python3.7/dist-packages/mujoco_py/generated/_pyxbld_2.0.2.13_37_linuxcpuextensionbuilder/temp.linux-x86_64-3.7/usr/local/lib/python3.7/dist-packages/mujoco_py/cymj.o -fopenmp -w\n",
            "x86_64-linux-gnu-gcc -pthread -Wno-unused-result -Wsign-compare -DNDEBUG -g -fwrapv -O2 -Wall -g -fstack-protector-strong -Wformat -Werror=format-security -g -fwrapv -O2 -g -fstack-protector-strong -Wformat -Werror=format-security -Wdate-time -D_FORTIFY_SOURCE=2 -fPIC -I/usr/local/lib/python3.7/dist-packages/mujoco_py -I/usr/local/lib/python3.7/dist-packages/mujoco_py/binaries/linux/mujoco210/include -I/usr/local/lib/python3.7/dist-packages/numpy/core/include -I/usr/include/python3.7m -c /usr/local/lib/python3.7/dist-packages/mujoco_py/gl/osmesashim.c -o /usr/local/lib/python3.7/dist-packages/mujoco_py/generated/_pyxbld_2.0.2.13_37_linuxcpuextensionbuilder/temp.linux-x86_64-3.7/usr/local/lib/python3.7/dist-packages/mujoco_py/gl/osmesashim.o -fopenmp -w\n",
            "creating /usr/local/lib/python3.7/dist-packages/mujoco_py/generated/_pyxbld_2.0.2.13_37_linuxcpuextensionbuilder/lib.linux-x86_64-3.7\n",
            "creating /usr/local/lib/python3.7/dist-packages/mujoco_py/generated/_pyxbld_2.0.2.13_37_linuxcpuextensionbuilder/lib.linux-x86_64-3.7/mujoco_py\n",
            "x86_64-linux-gnu-gcc -pthread -shared -Wl,-O1 -Wl,-Bsymbolic-functions -Wl,-Bsymbolic-functions -g -fwrapv -O2 -Wl,-Bsymbolic-functions -g -fwrapv -O2 -g -fstack-protector-strong -Wformat -Werror=format-security -Wdate-time -D_FORTIFY_SOURCE=2 /usr/local/lib/python3.7/dist-packages/mujoco_py/generated/_pyxbld_2.0.2.13_37_linuxcpuextensionbuilder/temp.linux-x86_64-3.7/usr/local/lib/python3.7/dist-packages/mujoco_py/cymj.o /usr/local/lib/python3.7/dist-packages/mujoco_py/generated/_pyxbld_2.0.2.13_37_linuxcpuextensionbuilder/temp.linux-x86_64-3.7/usr/local/lib/python3.7/dist-packages/mujoco_py/gl/osmesashim.o -L/usr/local/lib/python3.7/dist-packages/mujoco_py/binaries/linux/mujoco210/bin -Wl,--enable-new-dtags,-R/usr/local/lib/python3.7/dist-packages/mujoco_py/binaries/linux/mujoco210/bin -lmujoco210 -lglewosmesa -lOSMesa -lGL -o /usr/local/lib/python3.7/dist-packages/mujoco_py/generated/_pyxbld_2.0.2.13_37_linuxcpuextensionbuilder/lib.linux-x86_64-3.7/mujoco_py/cymj.cpython-37m-x86_64-linux-gnu.so -fopenmp\n"
          ]
        },
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "array([ 0.06148068,  0.02226833,  0.09254591,  0.09723181, -0.01016754,\n",
              "        0.05122158,  0.03591042, -0.05853889, -0.01213919, -0.03232064,\n",
              "        0.02659138,  0.17521858, -0.03403398, -0.07695705,  0.03207362,\n",
              "       -0.06039334,  0.03075459,  0.07048004,  0.        ,  0.7614807 ],\n",
              "      dtype=float32)"
            ]
          },
          "metadata": {},
          "execution_count": 4
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "EPSILON = 1e-6\n",
        "\n",
        "def linear_init(module):\n",
        "    if isinstance(module, nn.Linear):\n",
        "        nn.init.xavier_uniform_(module.weight)\n",
        "        module.bias.data.zero_()\n",
        "    return module\n",
        "\n",
        "\n",
        "class DiagNormalPolicy(nn.Module):\n",
        "\n",
        "    def __init__(self, input_size, output_size, hiddens=None, activation='relu', device='cpu'):\n",
        "        super(DiagNormalPolicy, self).__init__()\n",
        "        self.device = device\n",
        "        if hiddens is None:\n",
        "            hiddens = [100, 100]\n",
        "        if activation == 'relu':\n",
        "            activation = nn.ReLU\n",
        "        elif activation == 'tanh':\n",
        "            activation = nn.Tanh\n",
        "        layers = [linear_init(nn.Linear(input_size, hiddens[0])), activation()]\n",
        "        for i, o in zip(hiddens[:-1], hiddens[1:]):\n",
        "            layers.append(linear_init(nn.Linear(i, o)))\n",
        "            layers.append(activation())\n",
        "        layers.append(linear_init(nn.Linear(hiddens[-1], output_size)))\n",
        "        self.mean = nn.Sequential(*layers)\n",
        "        self.sigma = nn.Parameter(torch.Tensor(output_size))\n",
        "        self.sigma.data.fill_(math.log(1))\n",
        "\n",
        "    # def forward(self, state):\n",
        "    #     state = state.to(self.device, non_blocking=True)\n",
        "    #     loc = self.mean(state)\n",
        "    #     scale = torch.exp(torch.clamp(self.sigma, min=math.log(EPSILON)))\n",
        "    #     return Normal(loc=loc, scale=scale)\n",
        "\n",
        "    def density(self, state):\n",
        "        state = state.to(self.device, non_blocking=True)\n",
        "        loc = self.mean(state)\n",
        "        scale = torch.exp(torch.clamp(self.sigma, min=math.log(EPSILON)))\n",
        "        return Normal(loc=loc, scale=scale)\n",
        "\n",
        "    def log_prob(self, state, action):\n",
        "        density = self.density(state)\n",
        "        return density.log_prob(action).mean(dim=1, keepdim=True)\n",
        "\n",
        "    def forward(self, state):\n",
        "        density = self.density(state)\n",
        "        action = density.sample()\n",
        "        return action\n"
      ],
      "metadata": {
        "id": "bUMYPyEX3xz4"
      },
      "execution_count": 5,
      "outputs": []
    },
    {
      "cell_type": "code",
      "execution_count": 6,
      "metadata": {
        "id": "9x6TJOktOIcF"
      },
      "outputs": [],
      "source": [
        "def layer_init(layer, std=np.sqrt(2), bias_const=0.0):\n",
        "    torch.nn.init.orthogonal_(layer.weight, std)\n",
        "    torch.nn.init.constant_(layer.bias, bias_const)\n",
        "    return layer"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 7,
      "metadata": {
        "id": "MDEKGt5KONHM"
      },
      "outputs": [],
      "source": [
        "class Actor(nn.Module):\n",
        "    def __init__(self, env, lr, hidden_size=100):\n",
        "        super().__init__()\n",
        "        self.input_size = env.observation_space.shape[0]\n",
        "        self.actor_output_size = env.action_space.shape[0]\n",
        "\n",
        "        self.l1 = layer_init(nn.Linear(self.input_size, hidden_size))\n",
        "        self.l2 = layer_init(nn.Linear(hidden_size, hidden_size))\n",
        "        self.output = layer_init(nn.Linear(hidden_size, self.actor_output_size), std=0.01)\n",
        "        self.activation = nn.ReLU()\n",
        "        self.distribution = ch.distributions.ActionDistribution(env)\n",
        "\n",
        "        self.optimizer = torch.optim.Adam(self.parameters(), lr=lr, eps=1e-5)\n",
        "\n",
        "    def forward(self, x):\n",
        "        x = self.activation(self.l1(x))\n",
        "        x = self.activation(self.l2(x))\n",
        "        x = self.output(x)\n",
        "        mass = self.distribution(x)\n",
        "\n",
        "        return mass"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 8,
      "metadata": {
        "id": "Dbv0dy_WOOf9"
      },
      "outputs": [],
      "source": [
        "class Critic(nn.Module):\n",
        "    def __init__(self, env, lr=0.01, hidden_size=32):\n",
        "        super().__init__()\n",
        "        self.input_size = env.observation_space.shape[0]\n",
        "        self.critic_output_size = 1\n",
        "\n",
        "        self.l1 = layer_init(nn.Linear(self.input_size, hidden_size))\n",
        "        self.l2 = layer_init(nn.Linear(hidden_size, hidden_size))\n",
        "        self.critic_head = layer_init(nn.Linear(hidden_size, self.critic_output_size), std=1.)\n",
        "        self.activation = nn.ReLU()\n",
        "\n",
        "        self.optimizer = torch.optim.Adam(self.parameters(), lr=lr, eps=1e-5)\n",
        "\n",
        "    def forward(self, x):\n",
        "        x = self.activation(self.l1(x))\n",
        "        x = self.activation(self.l2(x))\n",
        "        value = self.critic_head(x)\n",
        "\n",
        "        return value"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 10,
      "metadata": {
        "id": "YIPEWKrndmIl"
      },
      "outputs": [],
      "source": [
        "class MAMLPPO():\n",
        "    def __init__(self, env_name,\n",
        "                 actor_class=Actor, critic_class=Critic, \n",
        "                 actor_args=dict(), critic_args=dict(),\n",
        "                 adapt_lr=1e-1, meta_lr=1e-2, \n",
        "                 adapt_steps=3, ppo_steps=5,\n",
        "                 adapt_batch_size=20, meta_batch_size=20,\n",
        "                 gamma=0.99, tau=1.0,\n",
        "                 policy_clip=0.2, value_clip=None,\n",
        "                 num_workers=10,\n",
        "                 seed=42,\n",
        "                 device=None, name=\"MAMLPPO\", tensorboard_log=\"./logs\"):\n",
        "        \n",
        "        random.seed(seed)\n",
        "        np.random.seed(seed)\n",
        "        torch.manual_seed(seed)\n",
        "        \n",
        "        if torch.cuda.is_available():\n",
        "            self.device = torch.device(\"cuda\")\n",
        "            torch.cuda.manual_seed(seed)\n",
        "        else:\n",
        "            self.device = torch.device(\"cpu\")\n",
        "        if device:\n",
        "            self.device = torch.device(device)\n",
        "        print(\"Running on: \" + str(self.device))\n",
        "\n",
        "        def make_env():\n",
        "            env = gym.make(env_name)\n",
        "            env = ch.envs.ActionSpaceScaler(env)\n",
        "            return env\n",
        "\n",
        "        env = l2l.gym.AsyncVectorEnv([make_env for _ in range(num_workers)])\n",
        "        env.seed(seed)\n",
        "        env.set_task(env.sample_tasks(1)[0])\n",
        "        self.env = ch.envs.Torch(env)\n",
        "\n",
        "        self.gamma = gamma\n",
        "        self.tau = tau\n",
        "        self.adapt_lr = adapt_lr\n",
        "        self.meta_lr = meta_lr\n",
        "        self.adapt_steps = adapt_steps\n",
        "        self.adapt_batch_size = adapt_batch_size\n",
        "        self.meta_batch_size = meta_batch_size\n",
        "        self.policy_clip = policy_clip\n",
        "        self.value_clip = value_clip\n",
        "        self.ppo_steps = ppo_steps\n",
        "        self.global_iteration = 0\n",
        "\n",
        "        # self.policy = Actor(env, **actor_args).to(device)\n",
        "        # self.baseline = Critic(env, lr=0.001, **critic_args).to(device)\n",
        "        self.policy = DiagNormalPolicy(self.env.state_size, self.env.action_size, device=self.device)\n",
        "        self.baseline = LinearValue(self.env.state_size, self.env.action_size)\n",
        "\n",
        "        self.policy.to(self.device)\n",
        "        self.baseline.to(self.device)\n",
        "\n",
        "        self.optimizer = torch.optim.Adam(self.policy.parameters(), meta_lr)\n",
        "\n",
        "        if tensorboard_log is not None:\n",
        "            self.run_name = name + \"_\" + str(int(time.time()))\n",
        "            self.writer = SummaryWriter(f\"{tensorboard_log}/{self.run_name}\")\n",
        "        else:\n",
        "            self.writer = None\n",
        "\n",
        "\n",
        "    def save(self, path=\"./\"):\n",
        "        torch.save(self.baseline.state_dict(), path + \"/baseline.pt\")\n",
        "        torch.save(self.policy.state_dict(), path + \"/policy.pt\")\n",
        "\n",
        "\n",
        "    def load(self, path=\"./\"):\n",
        "        self.baseline.load_state_dict(torch.load(path + \"/baseline.pt\"))\n",
        "        self.policy.load_state_dict(torch.load(path + \"/policy.pt\"))\n",
        "\n",
        "\n",
        "    def collect_steps(self, policy, n_episodes):\n",
        "        # replay = ch.ExperienceReplay(device=self.device)\n",
        "        # for i in range(n_episodes):\n",
        "        #     state = self.env.reset()\n",
        "\n",
        "        #     while True:\n",
        "        #         with torch.no_grad():\n",
        "        #             mass = policy.density(state)\n",
        "        #         action = mass.sample()\n",
        "        #         log_prob = mass.log_prob(action).mean(dim=1, keepdim=True)\n",
        "        #         next_state, reward, done, _ = self.env.step(action)\n",
        "\n",
        "                # replay.append(state,\n",
        "                #             action,\n",
        "                #             reward,\n",
        "                #             next_state,\n",
        "                #             done,\n",
        "                #             log_prob=log_prob)\n",
        "                \n",
        "        #         if done.any():\n",
        "        #             break\n",
        "\n",
        "        #         state = next_state\n",
        "\n",
        "        self.env.reset()\n",
        "        task = ch.envs.Runner(self.env)\n",
        "        replay = task.run(policy, episodes=n_episodes).to(self.device)\n",
        "\n",
        "        with torch.no_grad():\n",
        "            next_state_value = self.baseline(replay[-1].next_state)\n",
        "            mass = policy.density(replay.state())\n",
        "\n",
        "        log_probs = mass.log_prob(replay.action()).mean(dim=1, keepdim=True)\n",
        "        values = self.baseline(replay.state())\n",
        "\n",
        "        advantages = ch.generalized_advantage(self.gamma,\n",
        "                                                self.tau,\n",
        "                                                replay.reward(),\n",
        "                                                replay.done(),\n",
        "                                                values.detach(),\n",
        "                                                next_state_value)\n",
        "        returns = advantages + values.detach()\n",
        "        advantages = ch.normalize(advantages, epsilon=1e-8)\n",
        "\n",
        "        for i, sars in enumerate(replay):\n",
        "            sars.returns = returns[i]\n",
        "            sars.advantage = advantages[i]\n",
        "            sars.log_prob = log_probs[i]\n",
        "\n",
        "        # if self.value_clip:\n",
        "        #     value_loss = ppo.state_value_loss(values,\n",
        "        #                                     replay.value(),\n",
        "        #                                     returns,\n",
        "        #                                     clip=self.value_clip)\n",
        "        # else:\n",
        "        #     value_loss = a2c.state_value_loss(values, returns)\n",
        "        # self.baseline.optimizer.zero_grad()\n",
        "        # value_loss.backward()\n",
        "        # self.baseline.optimizer.step()\n",
        "\n",
        "        self.baseline.fit(replay.state(), returns)\n",
        "        return replay\n",
        "\n",
        "\n",
        "    def maml_a2c_loss(self, train_episodes, learner):\n",
        "        # Update policy and baseline\n",
        "        states = train_episodes.state()\n",
        "        actions = train_episodes.action()\n",
        "        density = learner.density(states)\n",
        "        log_probs = density.log_prob(actions).mean(dim=1, keepdim=True)\n",
        "\n",
        "        advantages = train_episodes.advantage()\n",
        "        return a2c.policy_loss(log_probs, train_episodes.advantage())\n",
        "\n",
        "\n",
        "    def fast_adapt(self, clone, train_episodes, first_order=False):\n",
        "        second_order = not first_order\n",
        "        loss = self.maml_a2c_loss(train_episodes, clone)\n",
        "        gradients = autograd.grad(loss,\n",
        "                                clone.parameters(),\n",
        "                                retain_graph=second_order,\n",
        "                                create_graph=second_order)\n",
        "        return l2l.algorithms.maml.maml_update(clone, self.adapt_lr, gradients)\n",
        "        \n",
        "\n",
        "    def meta_loss(self, iteration_replays, iteration_policies, policy):\n",
        "        mean_loss = 0.0\n",
        "        for task_replays, old_policy in tqdm(zip(iteration_replays, iteration_policies),\n",
        "                                            total=len(iteration_replays),\n",
        "                                            desc='Surrogate Loss',\n",
        "                                            leave=False):\n",
        "            train_replays = task_replays[:-1]\n",
        "            valid_episodes = task_replays[-1]\n",
        "            new_policy = l2l.clone_module(policy)\n",
        "\n",
        "            # Fast Adapt\n",
        "            for train_episodes in train_replays:\n",
        "                new_policy = self.fast_adapt(new_policy, train_episodes, first_order=False)\n",
        "\n",
        "            # Useful values\n",
        "            states = valid_episodes.state()\n",
        "            actions = valid_episodes.action()\n",
        "\n",
        "            # Compute KL\n",
        "            old_densities = old_policy.density(states)\n",
        "            new_densities = new_policy.density(states)\n",
        "\n",
        "            # Compute Surrogate Loss\n",
        "            advantages = valid_episodes.advantage()\n",
        "            old_log_probs = old_densities.log_prob(actions).mean(dim=1, keepdim=True).detach()\n",
        "            new_log_probs = new_densities.log_prob(actions).mean(dim=1, keepdim=True)\n",
        "            mean_loss += ppo.policy_loss(new_log_probs, \n",
        "                                         old_log_probs, \n",
        "                                         advantages,\n",
        "                                         clip=self.policy_clip)\n",
        "        mean_loss /= len(iteration_replays)\n",
        "        return mean_loss\n",
        "\n",
        "\n",
        "    def meta_optimize(self, iteration_replays, iteration_policies):\n",
        "        loss = self.meta_loss(iteration_replays, iteration_policies, self.policy)\n",
        "\n",
        "        self.optimizer.zero_grad()\n",
        "        loss.backward()\n",
        "        self.optimizer.step()\n",
        "        \n",
        "        if self.writer is not None:\n",
        "            self.writer.add_scalar(\"loss\", loss, self.global_iteration)\n",
        "\n",
        "\n",
        "    def train(self, num_iterations=100):\n",
        "        for iteration in range(num_iterations):\n",
        "            self.global_iteration += 1\n",
        "            iteration_reward = 0.0\n",
        "            iteration_replays = []\n",
        "            iteration_policies = []\n",
        "            iter_loss = 0.0\n",
        "\n",
        "            for task_config in tqdm(self.env.sample_tasks(self.meta_batch_size), leave=False, desc='Data'):\n",
        "                clone = deepcopy(self.policy)\n",
        "                self.env.set_task(task_config)\n",
        "                task_replay = []\n",
        "\n",
        "                # Fast Adapt\n",
        "                for step in range(self.adapt_steps):\n",
        "                    train_episodes = self.collect_steps(clone, n_episodes=self.adapt_batch_size)\n",
        "                    self.fast_adapt(clone, train_episodes, first_order=True)\n",
        "                    task_replay.append(train_episodes)\n",
        "\n",
        "                # Compute Validation Loss\n",
        "                valid_episodes = self.collect_steps(clone, n_episodes=self.adapt_batch_size)\n",
        "                task_replay.append(valid_episodes)\n",
        "                iteration_reward += valid_episodes.reward().sum().item() / self.adapt_batch_size\n",
        "                iteration_replays.append(task_replay)\n",
        "                iteration_policies.append(clone)\n",
        "\n",
        "            # Print statistics\n",
        "            print('\\nIteration', self.global_iteration)\n",
        "            adaptation_reward = iteration_reward / self.meta_batch_size\n",
        "            print('adaptation_reward', adaptation_reward)\n",
        "\n",
        "            if self.writer is not None:\n",
        "                self.writer.add_scalar(\"adaptation_reward\", adaptation_reward, self.global_iteration)\n",
        "\n",
        "            self.meta_optimize(iteration_replays, iteration_policies)\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 11,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "qhHiZBGTPkJZ",
        "outputId": "0b7a51d7-d5c6-48f1-c9ce-2288a902de47"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Running on: cuda\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "/usr/local/lib/python3.7/dist-packages/gym/logger.py:30: UserWarning: \u001b[33mWARN: Box bound precision lowered by casting to float32\u001b[0m\n",
            "  warnings.warn(colorize('%s: %s'%('WARN', msg % args), 'yellow'))\n"
          ]
        }
      ],
      "source": [
        "aa = MAMLPPO('Particles2D-v1')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 12,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 1000
        },
        "id": "1UAU7YG-Pn34",
        "outputId": "7a0529b0-dcbb-4a0a-f1c1-716860ce6a48"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            ""
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\n",
            "Iteration 1\n",
            "adaptation_reward -67.36041107177735\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            ""
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\n",
            "Iteration 2\n",
            "adaptation_reward -49.342748870849604\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            ""
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\n",
            "Iteration 3\n",
            "adaptation_reward -42.9727781677246\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            ""
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\n",
            "Iteration 4\n",
            "adaptation_reward -36.72748077392579\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            ""
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\n",
            "Iteration 5\n",
            "adaptation_reward -40.57419281005859\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            ""
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\n",
            "Iteration 6\n",
            "adaptation_reward -38.39564487457275\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            ""
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\n",
            "Iteration 7\n",
            "adaptation_reward -37.70770763397216\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            ""
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\n",
            "Iteration 8\n",
            "adaptation_reward -39.86008567810059\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            ""
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\n",
            "Iteration 9\n",
            "adaptation_reward -38.29249656677247\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            ""
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\n",
            "Iteration 10\n",
            "adaptation_reward -42.756781845092775\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            ""
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\n",
            "Iteration 11\n",
            "adaptation_reward -39.44821411132813\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            ""
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\n",
            "Iteration 12\n",
            "adaptation_reward -38.16800643920899\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            ""
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\n",
            "Iteration 13\n",
            "adaptation_reward -38.74898471832277\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            ""
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\n",
            "Iteration 14\n",
            "adaptation_reward -38.12704372406006\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            ""
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\n",
            "Iteration 15\n",
            "adaptation_reward -33.669792346954345\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            ""
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\n",
            "Iteration 16\n",
            "adaptation_reward -36.29836557388305\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            ""
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\n",
            "Iteration 17\n",
            "adaptation_reward -39.41265470504761\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            ""
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\n",
            "Iteration 18\n",
            "adaptation_reward -43.88270620346069\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            ""
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\n",
            "Iteration 19\n",
            "adaptation_reward -46.77267856597899\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            ""
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\n",
            "Iteration 20\n",
            "adaptation_reward -44.06627119064332\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            ""
          ]
        },
        {
          "output_type": "error",
          "ename": "KeyboardInterrupt",
          "evalue": "ignored",
          "traceback": [
            "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
            "\u001b[0;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
            "\u001b[0;32m<ipython-input-12-a3eb2274f928>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m()\u001b[0m\n\u001b[0;32m----> 1\u001b[0;31m \u001b[0maa\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtrain\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;36m100\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m",
            "\u001b[0;32m<ipython-input-10-079f75097aa5>\u001b[0m in \u001b[0;36mtrain\u001b[0;34m(self, num_iterations)\u001b[0m\n\u001b[1;32m    219\u001b[0m                 \u001b[0;31m# Fast Adapt\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    220\u001b[0m                 \u001b[0;32mfor\u001b[0m \u001b[0mstep\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mrange\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0madapt_steps\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 221\u001b[0;31m                     \u001b[0mtrain_episodes\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcollect_steps\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mclone\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mn_episodes\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0madapt_batch_size\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    222\u001b[0m                     \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mfast_adapt\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mclone\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtrain_episodes\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mfirst_order\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;32mTrue\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    223\u001b[0m                     \u001b[0mtask_replay\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mappend\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mtrain_episodes\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m<ipython-input-10-079f75097aa5>\u001b[0m in \u001b[0;36mcollect_steps\u001b[0;34m(self, policy, n_episodes)\u001b[0m\n\u001b[1;32m    100\u001b[0m         \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0menv\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mreset\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    101\u001b[0m         \u001b[0mtask\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0menvs\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mRunner\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0menv\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 102\u001b[0;31m         \u001b[0mreplay\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mtask\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mrun\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mpolicy\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mepisodes\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mn_episodes\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mto\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdevice\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    103\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    104\u001b[0m         \u001b[0;32mwith\u001b[0m \u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mno_grad\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.7/dist-packages/cherry/envs/runner_wrapper.py\u001b[0m in \u001b[0;36mrun\u001b[0;34m(self, get_action, steps, episodes, render)\u001b[0m\n\u001b[1;32m    115\u001b[0m                 \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mreset\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    116\u001b[0m             \u001b[0minfo\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m{\u001b[0m\u001b[0;34m}\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 117\u001b[0;31m             \u001b[0maction\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mget_action\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_current_state\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    118\u001b[0m             \u001b[0;32mif\u001b[0m \u001b[0misinstance\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0maction\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtuple\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    119\u001b[0m                 \u001b[0mskip_unpack\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;32mFalse\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.7/dist-packages/torch/nn/modules/module.py\u001b[0m in \u001b[0;36m_call_impl\u001b[0;34m(self, *input, **kwargs)\u001b[0m\n\u001b[1;32m   1128\u001b[0m         if not (self._backward_hooks or self._forward_hooks or self._forward_pre_hooks or _global_backward_hooks\n\u001b[1;32m   1129\u001b[0m                 or _global_forward_hooks or _global_forward_pre_hooks):\n\u001b[0;32m-> 1130\u001b[0;31m             \u001b[0;32mreturn\u001b[0m \u001b[0mforward_call\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0minput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1131\u001b[0m         \u001b[0;31m# Do not call functions when jit is used\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1132\u001b[0m         \u001b[0mfull_backward_hooks\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mnon_full_backward_hooks\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m<ipython-input-5-9d91b8135330>\u001b[0m in \u001b[0;36mforward\u001b[0;34m(self, state)\u001b[0m\n\u001b[1;32m     46\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0mforward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mstate\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     47\u001b[0m         \u001b[0mdensity\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdensity\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mstate\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 48\u001b[0;31m         \u001b[0maction\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mdensity\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0msample\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     49\u001b[0m         \u001b[0;32mreturn\u001b[0m \u001b[0maction\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.7/dist-packages/torch/distributions/normal.py\u001b[0m in \u001b[0;36msample\u001b[0;34m(self, sample_shape)\u001b[0m\n\u001b[1;32m     66\u001b[0m         \u001b[0mshape\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_extended_shape\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0msample_shape\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     67\u001b[0m         \u001b[0;32mwith\u001b[0m \u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mno_grad\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 68\u001b[0;31m             \u001b[0;32mreturn\u001b[0m \u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mnormal\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mloc\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mexpand\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mshape\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mscale\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mexpand\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mshape\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     69\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     70\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0mrsample\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0msample_shape\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mSize\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;31mKeyboardInterrupt\u001b[0m: "
          ]
        }
      ],
      "source": [
        "aa.train(100)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "iYFpGEDjPwXy"
      },
      "outputs": [],
      "source": [
        ""
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "collapsed_sections": [],
      "name": "MAML A2C.ipynb",
      "provenance": []
    },
    "gpuClass": "standard",
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}