MAML_PPO.ipynb

{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "0u8jyVEHU90_",
        "outputId": "95d2415f-493d-427e-a10b-36d7499b002c"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Reading package lists... Done\n",
            "Building dependency tree       \n",
            "Reading state information... Done\n",
            "libglew-dev is already the newest version (2.0.0-5).\n",
            "libgl1-mesa-dev is already the newest version (20.0.8-0ubuntu1~18.04.1).\n",
            "libgl1-mesa-glx is already the newest version (20.0.8-0ubuntu1~18.04.1).\n",
            "libosmesa6-dev is already the newest version (20.0.8-0ubuntu1~18.04.1).\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",
            "0 upgraded, 0 newly installed, 0 to remove and 19 not upgraded.\n",
            "Reading package lists... Done\n",
            "Building dependency tree       \n",
            "Reading state information... Done\n",
            "patchelf is already the newest version (0.9-1).\n",
            "The following package was automatically installed and is no longer required:\n",
            "  libnvidia-common-460\n",
            "Use 'apt autoremove' to remove it.\n",
            "0 upgraded, 0 newly installed, 0 to remove and 19 not upgraded.\n",
            "Looking in indexes: https://pypi.org/simple, https://us-python.pkg.dev/colab-wheels/public/simple/\n",
            "Requirement already satisfied: free-mujoco-py in /usr/local/lib/python3.7/dist-packages (2.1.6)\n",
            "Requirement 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: 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",
            "Requirement already satisfied: glfw<2.0.0,>=1.4.0 in /usr/local/lib/python3.7/dist-packages (from free-mujoco-py) (1.12.0)\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: fasteners==0.15 in /usr/local/lib/python3.7/dist-packages (from free-mujoco-py) (0.15)\n",
            "Requirement already satisfied: monotonic>=0.1 in /usr/local/lib/python3.7/dist-packages (from fasteners==0.15->free-mujoco-py) (1.6)\n",
            "Requirement already satisfied: six in /usr/local/lib/python3.7/dist-packages (from fasteners==0.15->free-mujoco-py) (1.15.0)\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"
          ]
        }
      ],
      "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": null,
      "metadata": {
        "id": "wZleUELMTDCf"
      },
      "outputs": [],
      "source": [
        "!pip install cherry-rl learn2learn &> /dev/null"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "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": "f705e7b2-14e2-4c3d-d2ef-e2955e75bd40"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "array([ 0.00295027,  0.00932332, -0.09128403, -0.06178263,  0.07330716,\n",
              "       -0.09275652,  0.02119905, -0.03270184,  0.11129588, -0.0150656 ,\n",
              "       -0.05936605,  0.03230394,  0.06749459,  0.17624147, -0.08633809,\n",
              "        0.081727  ,  0.01158038,  0.03684237,  0.        ,  0.7029503 ],\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": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "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": null,
      "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": null,
      "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": null,
      "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",
        "        for ppo_epoch in range(self.ppo_steps):\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": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "qhHiZBGTPkJZ",
        "outputId": "63219c2d-a4e8-4017-8cb4-7a0710d50423"
      },
      "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('HalfCheetahForwardBackward-v1')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "1UAU7YG-Pn34",
        "outputId": "43f2da3b-4d29-4546-b5ae-d35ce86667b6"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\rData:   0%|          | 0/20 [00:00<?, ?it/s]/usr/local/lib/python3.7/dist-packages/cherry/models/robotics.py:159: UserWarning: torch.lstsq is deprecated in favor of torch.linalg.lstsq and will be removed in a future PyTorch release.\n",
            "torch.linalg.lstsq has reversed arguments and does not return the QR decomposition in the returned tuple (although it returns other information about the problem).\n",
            "To get the qr decomposition consider using torch.linalg.qr.\n",
            "The returned solution in torch.lstsq stored the residuals of the solution in the last m - n columns of the returned value whenever m > n. In torch.linalg.lstsq, the residuals in the field 'residuals' of the returned named tuple.\n",
            "The unpacking of the solution, as in\n",
            "X, _ = torch.lstsq(B, A).solution[:A.size(1)]\n",
            "should be replaced with\n",
            "X = torch.linalg.lstsq(A, B).solution (Triggered internally at  ../aten/src/ATen/native/cuda/linalg/BatchLinearAlgebra.cpp:3276.)\n",
            "  coeffs, _ = th.lstsq(b, A)\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\n",
            "Iteration 1\n",
            "adaptation_reward -15.239832963943481\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            ""
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\n",
            "Iteration 2\n",
            "adaptation_reward -21.80439588546753\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            ""
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\n",
            "Iteration 3\n",
            "adaptation_reward -14.98247666358948\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            ""
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\n",
            "Iteration 4\n",
            "adaptation_reward -26.02806373596191\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            ""
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\n",
            "Iteration 5\n",
            "adaptation_reward -40.9422484588623\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            ""
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\n",
            "Iteration 6\n",
            "adaptation_reward -21.61346605300903\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            ""
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\n",
            "Iteration 7\n",
            "adaptation_reward -22.95705913543701\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            ""
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\n",
            "Iteration 8\n",
            "adaptation_reward -19.712455844879152\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            ""
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\n",
            "Iteration 9\n",
            "adaptation_reward -19.058130903244017\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            ""
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\n",
            "Iteration 10\n",
            "adaptation_reward -16.015074186325073\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            ""
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\n",
            "Iteration 11\n",
            "adaptation_reward -19.549395561218265\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            ""
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\n",
            "Iteration 12\n",
            "adaptation_reward -25.23745170593262\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            ""
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\n",
            "Iteration 13\n",
            "adaptation_reward -16.58354446411133\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            ""
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\n",
            "Iteration 14\n",
            "adaptation_reward -20.103928909301754\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            ""
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\n",
            "Iteration 15\n",
            "adaptation_reward -30.201771316528323\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            ""
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\n",
            "Iteration 16\n",
            "adaptation_reward -21.809322128295896\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            ""
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\n",
            "Iteration 17\n",
            "adaptation_reward -20.62111457824707\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            ""
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\n",
            "Iteration 18\n",
            "adaptation_reward -18.594097137451172\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            ""
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\n",
            "Iteration 19\n",
            "adaptation_reward -15.111018524169916\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            ""
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\n",
            "Iteration 20\n",
            "adaptation_reward -53.71457153320313\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            ""
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\n",
            "Iteration 21\n",
            "adaptation_reward -23.496156616210936\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            ""
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\n",
            "Iteration 22\n",
            "adaptation_reward -32.101380615234376\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            ""
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\n",
            "Iteration 23\n",
            "adaptation_reward -20.220051879882813\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            ""
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\n",
            "Iteration 24\n",
            "adaptation_reward -22.165815582275393\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            ""
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\n",
            "Iteration 25\n",
            "adaptation_reward -8.222628173828124\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            ""
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\n",
            "Iteration 26\n",
            "adaptation_reward -37.30268112182617\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            ""
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\n",
            "Iteration 27\n",
            "adaptation_reward -14.722551269531252\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            ""
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\n",
            "Iteration 28\n",
            "adaptation_reward -32.08169250488281\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            ""
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\n",
            "Iteration 29\n",
            "adaptation_reward -33.24332565307617\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            ""
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\n",
            "Iteration 30\n",
            "adaptation_reward -33.2961939239502\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            ""
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\n",
            "Iteration 31\n",
            "adaptation_reward -22.374862670898438\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            ""
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\n",
            "Iteration 32\n",
            "adaptation_reward -22.475312881469726\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            ""
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\n",
            "Iteration 33\n",
            "adaptation_reward -9.784945907592775\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            ""
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\n",
            "Iteration 34\n",
            "adaptation_reward -44.5146508026123\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            ""
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\n",
            "Iteration 35\n",
            "adaptation_reward -37.16360359191894\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            ""
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\n",
            "Iteration 36\n",
            "adaptation_reward -16.57937843322754\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            ""
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\n",
            "Iteration 37\n",
            "adaptation_reward -15.947536010742189\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            ""
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\n",
            "Iteration 38\n",
            "adaptation_reward -23.439432907104496\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            ""
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\n",
            "Iteration 39\n",
            "adaptation_reward -25.786630191802978\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            ""
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\n",
            "Iteration 40\n",
            "adaptation_reward -4.652285308837891\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            ""
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\n",
            "Iteration 41\n",
            "adaptation_reward -16.95811584472656\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            ""
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\n",
            "Iteration 42\n",
            "adaptation_reward -13.229454650878903\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            ""
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\n",
            "Iteration 43\n",
            "adaptation_reward -9.597256927490236\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            ""
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\n",
            "Iteration 44\n",
            "adaptation_reward -14.462917175292969\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            ""
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\n",
            "Iteration 45\n",
            "adaptation_reward -15.42728607177734\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            ""
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\n",
            "Iteration 46\n",
            "adaptation_reward -24.236006469726554\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            ""
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\n",
            "Iteration 47\n",
            "adaptation_reward -16.12834014892578\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            ""
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\n",
            "Iteration 48\n",
            "adaptation_reward -4.841907653808594\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            ""
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\n",
            "Iteration 49\n",
            "adaptation_reward -65.41198593139649\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            ""
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\n",
            "Iteration 50\n",
            "adaptation_reward -13.75320785522461\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            ""
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\n",
            "Iteration 51\n",
            "adaptation_reward -20.569122009277343\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            ""
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\n",
            "Iteration 52\n",
            "adaptation_reward -0.9126600646972678\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "Data:  20%|██        | 4/20 [00:21<01:25,  5.37s/it]"
          ]
        }
      ],
      "source": [
        "aa.train(100)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "iYFpGEDjPwXy"
      },
      "outputs": [],
      "source": [
        ""
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "collapsed_sections": [],
      "name": "New MAML PPO.ipynb",
      "provenance": []
    },
    "gpuClass": "standard",
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}