CP Predictor.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "complete-dating",
   "metadata": {},
   "source": [
    "# CP Predictor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "adequate-smoke",
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import Callable, Iterable\n",
    "import functools\n",
    "import numpy as np\n",
    "import os\n",
    "\n",
    "import gym\n",
    "\n",
    "import torch\n",
    "# torch.multiprocessing.set_start_method('spawn')\n",
    "import torch.nn as nn\n",
    "from torch.utils.data import IterableDataset, DataLoader\n",
    "from torch import optim\n",
    "\n",
    "from a2c_ppo_acktr import algo, utils\n",
    "from a2c_ppo_acktr.envs import make_vec_envs"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "convenient-sauce",
   "metadata": {},
   "source": [
    "## Predictors"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "insured-whole",
   "metadata": {},
   "outputs": [],
   "source": [
    "class PredictorMLP(nn.Module):\n",
    "    def __init__(self, in_size=128, hidden_size=256, out_size=128):\n",
    "        super().__init__()\n",
    "        self.fc_grp1 = nn.Sequential(\n",
    "            nn.Linear(in_size, hidden_size),\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(hidden_size, hidden_size),\n",
    "            nn.ReLU(),\n",
    "        )\n",
    "\n",
    "        self.action_fc = nn.Sequential(nn.Linear(1, hidden_size), nn.ReLU())\n",
    "\n",
    "        self.fc_grp2 = nn.Sequential(\n",
    "            nn.Linear(hidden_size, hidden_size),\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(hidden_size, out_size),\n",
    "        )\n",
    "\n",
    "    def forward(self, obs, action):\n",
    "        action = action.squeeze(-1)\n",
    "        obs_encoding = self.fc_grp1(obs)\n",
    "        action_encoding = self.action_fc(action)\n",
    "        prod = obs_encoding * action_encoding\n",
    "        out = self.fc_grp2(prod)\n",
    "        return out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "juvenile-studio",
   "metadata": {},
   "outputs": [],
   "source": [
    "class PredictorMLPTwo(nn.Module):\n",
    "    def __init__(self, in_size=128, hidden_size=256, out_size=128):\n",
    "        super().__init__()\n",
    "        self.fc_grp1 = nn.Sequential(\n",
    "            nn.Linear(in_size, hidden_size),\n",
    "            nn.ReLU(),\n",
    "            nn.BatchNorm1d(hidden_size),\n",
    "            nn.Dropout(p=0.4),\n",
    "            nn.Linear(hidden_size, hidden_size),\n",
    "            nn.ReLU(),\n",
    "            nn.BatchNorm1d(hidden_size),\n",
    "            nn.Dropout(p=0.4),\n",
    "            nn.Linear(hidden_size, hidden_size),\n",
    "            nn.ReLU(),\n",
    "        )\n",
    "\n",
    "        self.action_fc = nn.Sequential(nn.Linear(1, hidden_size), nn.ReLU())\n",
    "\n",
    "        self.fc_grp2 = nn.Sequential(\n",
    "            nn.Linear(hidden_size, hidden_size),\n",
    "            nn.ReLU(),\n",
    "            nn.BatchNorm1d(hidden_size),\n",
    "            nn.Dropout(p=0.4),\n",
    "            nn.Linear(hidden_size, hidden_size),\n",
    "            nn.ReLU(),\n",
    "            nn.BatchNorm1d(hidden_size),\n",
    "            nn.Dropout(p=0.4),\n",
    "            nn.Linear(hidden_size, out_size),\n",
    "        )\n",
    "\n",
    "    def forward(self, obs, action):\n",
    "        action = action.squeeze(-1)\n",
    "        obs_encoding = self.fc_grp1(obs)\n",
    "        action_encoding = self.action_fc(action)\n",
    "        prod = obs_encoding * action_encoding\n",
    "        out = self.fc_grp2(prod)\n",
    "        return out"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "promising-trace",
   "metadata": {},
   "source": [
    "## Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "transparent-extent",
   "metadata": {},
   "outputs": [],
   "source": [
    "class GeneratorDataset(IterableDataset):\n",
    "    \"\"\"Uses a generator function to generate batches of data.\"\"\"\n",
    "\n",
    "    def __init__(self, generator_fn: Callable):\n",
    "        self.generator_fn = generator_fn\n",
    "\n",
    "    def __iter__(self) -> Iterable:\n",
    "        return self.generator_fn()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "paperback-reduction",
   "metadata": {},
   "outputs": [],
   "source": [
    "def experience_gen(env, agent):\n",
    "    vis_obs = env.reset()\n",
    "    prev_obs = env.get_attr(\"unwrapped\")[0]._get_ram()\n",
    "    \n",
    "    if np.random.rand() < 0.1:\n",
    "        action = torch.randint(0, 6, (1,)).unsqueeze(-1)\n",
    "    else:\n",
    "        _, action, _, _ = agent.act(vis_obs, None, None)\n",
    "        action = action.cpu()\n",
    "\n",
    "    while True:\n",
    "        if np.random.rand() < 0.1:\n",
    "            action = torch.randint(0, 6, (1,)).unsqueeze(-1)\n",
    "        else:\n",
    "            _, action, _, _ = agent.act(vis_obs, None, None)\n",
    "            action = action.cpu()\n",
    "        \n",
    "        vis_obs, reward, done, _ = env.step(action)\n",
    "        cur_obs = env.get_attr(\"unwrapped\")[0]._get_ram()\n",
    "        if done:\n",
    "            yield prev_obs, action, cur_obs\n",
    "            vis_obs = env.reset()\n",
    "            prev_obs = env.get_attr(\"unwrapped\")[0]._get_ram()\n",
    "        else:\n",
    "            yield prev_obs, action, cur_obs\n",
    "            prev_obs = cur_obs"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dietary-stocks",
   "metadata": {},
   "source": [
    "## Training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "hairy-fault",
   "metadata": {},
   "outputs": [],
   "source": [
    "env_name = 'PongNoFrameskip-v4'\n",
    "seed = 1\n",
    "num_procs = 1\n",
    "gamma = 0.99\n",
    "log_dir = '/tmp/gym'\n",
    "device = 'cuda:0'\n",
    "envs = make_vec_envs(env_name, seed, num_procs,\n",
    "                     gamma, log_dir, device, False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "amazing-heaven",
   "metadata": {},
   "outputs": [],
   "source": [
    "trained_agent, _ = torch.load(\n",
    "            os.path.join(\"trained_models/a2c\", env_name + \".pt\"), map_location=device\n",
    "        )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "about-ridge",
   "metadata": {},
   "outputs": [],
   "source": [
    "trained_agent = trained_agent.eval()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "convinced-seven",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(4, 84, 84)"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "envs.observation_space.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "played-vehicle",
   "metadata": {},
   "outputs": [],
   "source": [
    "batch_size = 64"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "norwegian-organ",
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset = GeneratorDataset(functools.partial(experience_gen, env=envs, agent=trained_agent))\n",
    "dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=False, num_workers=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "guilty-suggestion",
   "metadata": {},
   "outputs": [
    {
     "ename": "NameError",
     "evalue": "name 'best_loss' is not defined",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mNameError\u001b[0m                                 Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-13-9e4783f8b59f>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[0;32m----> 1\u001b[0;31m \u001b[0mprint\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mbest_loss\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m",
      "\u001b[0;31mNameError\u001b[0m: name 'best_loss' is not defined"
     ]
    }
   ],
   "source": [
    "print(best_loss)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "id": "final-virtue",
   "metadata": {},
   "outputs": [],
   "source": [
    "lr = 1e-3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "id": "closed-lighting",
   "metadata": {},
   "outputs": [],
   "source": [
    "mask = torch.zeros(batch_size, 128)\n",
    "mask[:, 13] = 1\n",
    "mask[:, 14] = 1\n",
    "mask[:, 21] = 1\n",
    "mask[:, 51] = 2\n",
    "mask[:, 49] = 5\n",
    "mask[:, 54] = 5\n",
    "mask = mask.to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "id": "inappropriate-mouse",
   "metadata": {},
   "outputs": [],
   "source": [
    "def RAMLoss(pred, target):\n",
    "    loss = nn.MSELoss()\n",
    "    mse = loss(pred, target)\n",
    "    \n",
    "    ram_loss = (pred - target) * mask\n",
    "    ram_loss = ram_loss.pow(2).mean()\n",
    "#     print(ram_loss.shape)\n",
    "    \n",
    "#     cpu_score_loss = (pred[:, 13] - target[:, 13]).pow(2).mean()\n",
    "#     player_score_loss = (pred[:, 14] - target[:, 14]).pow(2).mean()\n",
    "#     cpu_paddle_loss = (pred[:, 21] - target[:, 21]).pow(2).mean()\n",
    "#     player_paddle_loss = (pred[:, 51] - target[:, 51]).pow(2).mean()\n",
    "#     ball_pos_loss = (pred[:, 49] - target[:, 49]).pow(2).mean() + (pred[:, 54] - target[:, 54]).pow(2).mean()\n",
    "    \n",
    "    return mse + ram_loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "id": "institutional-quality",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor(1086.2400, device='cuda:0', grad_fn=<AddBackward0>)"
      ]
     },
     "execution_count": 45,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "RAMLoss(pred, cur)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "id": "technical-delhi",
   "metadata": {},
   "outputs": [],
   "source": [
    "try:\n",
    "    net, best_loss = torch.load('trained_models/cp/predictor.pt', map_location=device)\n",
    "except (FileNotFoundError, TypeError):\n",
    "    net, best_loss = PredictorMLPTwo().to(device), 1000000"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "id": "linear-landscape",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "322.7232971191406\n"
     ]
    }
   ],
   "source": [
    "print(best_loss)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "id": "registered-habitat",
   "metadata": {},
   "outputs": [],
   "source": [
    "optimizer = optim.Adam(net.parameters(), lr=lr)\n",
    "# scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "id": "anticipated-wireless",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_lr(optimizer):\n",
    "    for param_group in optimizer.param_groups:\n",
    "        return param_group['lr']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "id": "excessive-representation",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0, 910.559814453125\n",
      "100, 1328.6968994140625\n",
      "200, 918.6124877929688\n",
      "300, 1159.4410400390625\n",
      "400, 891.9005126953125\n",
      "500, 1536.2138671875\n",
      "600, 974.6658935546875\n",
      "700, 700.4344482421875\n",
      "800, 1389.30908203125\n",
      "900, 1273.97705078125\n",
      "1000, 738.9677124023438\n",
      "1100, 927.365966796875\n",
      "1200, 815.876708984375\n",
      "1300, 1796.7689208984375\n",
      "1400, 809.6547241210938\n",
      "1500, 752.522216796875\n",
      "1600, 811.7977294921875\n",
      "1700, 1351.9534912109375\n",
      "1800, 667.17822265625\n",
      "1900, 850.1029052734375\n",
      "2000, 998.177490234375\n",
      "2100, 640.401123046875\n",
      "2200, 2020.424072265625\n",
      "2300, 758.96533203125\n",
      "2400, 1732.5469970703125\n",
      "2500, 822.7368774414062\n",
      "2600, 635.6709594726562\n",
      "2700, 790.884033203125\n",
      "2800, 1797.23828125\n",
      "2900, 1147.185791015625\n",
      "3000, 1233.4232177734375\n",
      "3100, 911.3038330078125\n",
      "3200, 926.6458740234375\n",
      "3300, 627.89990234375\n",
      "3400, 559.9838256835938\n",
      "3500, 1088.983154296875\n",
      "3600, 1068.4588623046875\n",
      "3700, 1086.847900390625\n",
      "3800, 533.4752807617188\n",
      "3900, 564.2099609375\n",
      "4000, 822.1962890625\n",
      "4100, 903.390625\n",
      "4200, 726.7750244140625\n",
      "4300, 1016.6827392578125\n",
      "4400, 670.9302368164062\n",
      "4500, 578.8619384765625\n",
      "4600, 709.2740478515625\n",
      "4700, 808.6275634765625\n",
      "4800, 795.6412353515625\n",
      "4900, 728.1947021484375\n",
      "5000, 663.1021728515625\n",
      "5100, 522.7653198242188\n",
      "5200, 830.8994140625\n",
      "5300, 1038.76220703125\n",
      "5400, 1341.56396484375\n",
      "5500, 742.712890625\n",
      "5600, 1562.072998046875\n",
      "5700, 687.6773681640625\n",
      "5800, 763.0032958984375\n",
      "5900, 657.248291015625\n",
      "6000, 408.30322265625\n",
      "6100, 894.9183349609375\n",
      "6200, 971.1234741210938\n",
      "6300, 719.2210693359375\n",
      "6400, 827.5550537109375\n",
      "6500, 578.156494140625\n",
      "6600, 1040.15087890625\n",
      "6700, 1087.6434326171875\n",
      "6800, 737.3701171875\n",
      "6900, 686.605712890625\n",
      "7000, 505.62396240234375\n",
      "7100, 777.611328125\n",
      "7200, 485.09429931640625\n",
      "7300, 604.037353515625\n",
      "7400, 600.4110717773438\n",
      "7500, 572.40478515625\n",
      "7600, 1057.5093994140625\n",
      "7700, 1223.095703125\n",
      "7800, 494.88751220703125\n",
      "7900, 494.57611083984375\n",
      "8000, 696.5057373046875\n",
      "8100, 2170.119873046875\n",
      "8200, 705.78369140625\n",
      "8300, 1835.01953125\n",
      "8400, 503.521728515625\n",
      "8500, 644.8660278320312\n",
      "8600, 422.82244873046875\n",
      "8700, 562.9051513671875\n",
      "8800, 573.1649169921875\n",
      "8900, 872.187255859375\n",
      "9000, 956.6304931640625\n",
      "9100, 520.0942993164062\n",
      "9200, 534.4631958007812\n",
      "9300, 569.8870849609375\n",
      "9400, 525.4027099609375\n",
      "9500, 1609.6854248046875\n",
      "9600, 555.1236572265625\n",
      "9700, 905.7607421875\n",
      "9800, 635.8076782226562\n",
      "9900, 481.2386474609375\n",
      "10000, 536.7614135742188\n",
      "10100, 1345.1832275390625\n",
      "10200, 565.5513916015625\n",
      "10300, 460.48370361328125\n",
      "10400, 737.2225952148438\n",
      "10500, 615.8363037109375\n",
      "10600, 1169.46826171875\n",
      "10700, 404.6559753417969\n",
      "10800, 610.2215576171875\n",
      "10900, 771.427001953125\n",
      "11000, 603.3255615234375\n",
      "11100, 774.814453125\n",
      "11200, 882.9427490234375\n",
      "11300, 559.3475341796875\n",
      "11400, 488.11572265625\n",
      "11500, 1083.1806640625\n",
      "11600, 667.1287841796875\n",
      "11700, 581.9583740234375\n",
      "11800, 607.9869995117188\n",
      "11900, 497.0281066894531\n",
      "12000, 495.58587646484375\n",
      "12100, 1117.97265625\n",
      "12200, 727.47705078125\n",
      "12300, 1605.77734375\n",
      "12400, 871.3203125\n",
      "12500, 778.3519287109375\n",
      "12600, 550.8076171875\n",
      "12700, 697.4390258789062\n",
      "12800, 500.5579833984375\n",
      "12900, 702.2091674804688\n",
      "13000, 595.5196533203125\n",
      "13100, 679.653076171875\n",
      "13200, 1545.7086181640625\n",
      "13300, 576.8773193359375\n",
      "13400, 840.4127807617188\n",
      "13500, 2037.311767578125\n",
      "13600, 829.1431884765625\n",
      "13700, 628.0521240234375\n",
      "13800, 2142.99755859375\n",
      "13900, 612.4188232421875\n",
      "14000, 621.6300659179688\n",
      "14100, 845.6926879882812\n",
      "14200, 741.9493408203125\n",
      "14300, 1393.524658203125\n",
      "14400, 647.273681640625\n",
      "14500, 1340.630615234375\n",
      "14600, 1243.268798828125\n",
      "14700, 781.6984252929688\n",
      "14800, 466.361328125\n",
      "14900, 440.8704833984375\n",
      "15000, 787.7867431640625\n",
      "15100, 620.5817260742188\n",
      "15200, 648.6739501953125\n",
      "15300, 378.21234130859375\n",
      "15400, 532.2095947265625\n",
      "15500, 739.7174682617188\n",
      "15600, 350.6260986328125\n",
      "15700, 430.3560485839844\n",
      "15800, 436.56707763671875\n",
      "15900, 506.3638916015625\n",
      "16000, 667.3098754882812\n",
      "16100, 950.5179443359375\n",
      "16200, 592.6575927734375\n",
      "16300, 692.2520141601562\n",
      "16400, 523.1251220703125\n",
      "16500, 466.67523193359375\n",
      "16600, 905.0623168945312\n",
      "16700, 631.3760986328125\n",
      "16800, 677.5014038085938\n",
      "16900, 824.316650390625\n",
      "17000, 963.2946166992188\n",
      "17100, 770.2891845703125\n",
      "17200, 1011.3672485351562\n",
      "17300, 1032.06005859375\n",
      "17400, 508.9921875\n",
      "17500, 679.6465454101562\n",
      "17600, 700.9769287109375\n",
      "17700, 881.056884765625\n",
      "17800, 752.5064697265625\n",
      "17900, 670.4899291992188\n",
      "18000, 658.5909423828125\n",
      "18100, 1183.4224853515625\n",
      "18200, 579.7392578125\n",
      "18300, 482.0882568359375\n",
      "18400, 813.2034301757812\n",
      "18500, 657.0814208984375\n",
      "18600, 368.62957763671875\n",
      "18700, 712.4671630859375\n",
      "18800, 681.5623168945312\n",
      "18900, 1199.764892578125\n",
      "19000, 1119.217041015625\n",
      "19100, 703.3491821289062\n",
      "19200, 1183.6287841796875\n"
     ]
    },
    {
     "ename": "KeyboardInterrupt",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-50-f962f2f68857>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[0;32m----> 1\u001b[0;31m \u001b[0;32mfor\u001b[0m \u001b[0mi\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m(\u001b[0m\u001b[0mprev\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0maction\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mcur\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mzip\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mrange\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;36m20000\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mdataloader\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      2\u001b[0m     \u001b[0mprev\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mprev\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mfloat\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mto\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[1;32m      3\u001b[0m     \u001b[0maction\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0maction\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mfloat\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mto\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[1;32m      4\u001b[0m     \u001b[0mcur\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mcur\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mfloat\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mto\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[1;32m      5\u001b[0m     \u001b[0mpred\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mnet\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mprev\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0maction\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/miniconda3/envs/ikos/lib/python3.8/site-packages/torch/utils/data/dataloader.py\u001b[0m in \u001b[0;36m__next__\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m    433\u001b[0m         \u001b[0;32mif\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_sampler_iter\u001b[0m \u001b[0;32mis\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    434\u001b[0m             \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_reset\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--> 435\u001b[0;31m         \u001b[0mdata\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_next_data\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    436\u001b[0m         \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_num_yielded\u001b[0m \u001b[0;34m+=\u001b[0m \u001b[0;36m1\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    437\u001b[0m         \u001b[0;32mif\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_dataset_kind\u001b[0m \u001b[0;34m==\u001b[0m \u001b[0m_DatasetKind\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mIterable\u001b[0m \u001b[0;32mand\u001b[0m\u001b[0;31m \u001b[0m\u001b[0;31m\\\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/miniconda3/envs/ikos/lib/python3.8/site-packages/torch/utils/data/dataloader.py\u001b[0m in \u001b[0;36m_next_data\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m    473\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0m_next_data\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\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    474\u001b[0m         \u001b[0mindex\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_next_index\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m  \u001b[0;31m# may raise StopIteration\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 475\u001b[0;31m         \u001b[0mdata\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_dataset_fetcher\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mfetch\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mindex\u001b[0m\u001b[0;34m)\u001b[0m  \u001b[0;31m# may raise StopIteration\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    476\u001b[0m         \u001b[0;32mif\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_pin_memory\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    477\u001b[0m             \u001b[0mdata\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0m_utils\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mpin_memory\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mpin_memory\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mdata\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/miniconda3/envs/ikos/lib/python3.8/site-packages/torch/utils/data/_utils/fetch.py\u001b[0m in \u001b[0;36mfetch\u001b[0;34m(self, possibly_batched_index)\u001b[0m\n\u001b[1;32m     26\u001b[0m             \u001b[0;32mfor\u001b[0m \u001b[0m_\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mpossibly_batched_index\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     27\u001b[0m                 \u001b[0;32mtry\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 28\u001b[0;31m                     \u001b[0mdata\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mappend\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mnext\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdataset_iter\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     29\u001b[0m                 \u001b[0;32mexcept\u001b[0m \u001b[0mStopIteration\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     30\u001b[0m                     \u001b[0;32mbreak\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m<ipython-input-4-1c49294b87d6>\u001b[0m in \u001b[0;36mexperience_gen\u001b[0;34m(env, agent)\u001b[0m\n\u001b[1;32m     16\u001b[0m             \u001b[0maction\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0maction\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcpu\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     17\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 18\u001b[0;31m         \u001b[0mvis_obs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mreward\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mdone\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0m_\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0menv\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mstep\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0maction\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     19\u001b[0m         \u001b[0mcur_obs\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0menv\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mget_attr\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m\"unwrapped\"\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;36m0\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_get_ram\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     20\u001b[0m         \u001b[0;32mif\u001b[0m \u001b[0mdone\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/miniconda3/envs/ikos/lib/python3.8/site-packages/stable_baselines3/common/vec_env/base_vec_env.py\u001b[0m in \u001b[0;36mstep\u001b[0;34m(self, actions)\u001b[0m\n\u001b[1;32m    161\u001b[0m         \"\"\"\n\u001b[1;32m    162\u001b[0m         \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mstep_async\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mactions\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 163\u001b[0;31m         \u001b[0;32mreturn\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mstep_wait\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    164\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    165\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0mget_images\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;34m->\u001b[0m \u001b[0mSequence\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mnp\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mndarray\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~/Courses/UCSD/CSE257/pytorch-a2c-ppo-acktr-gail/a2c_ppo_acktr/envs.py\u001b[0m in \u001b[0;36mstep_wait\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m    238\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    239\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0mstep_wait\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\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--> 240\u001b[0;31m         \u001b[0mobs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mrews\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mnews\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0minfos\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mvenv\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mstep_wait\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    241\u001b[0m         \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mstacked_obs\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[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mshape_dim0\u001b[0m\u001b[0;34m]\u001b[0m \u001b[0;34m=\u001b[0m\u001b[0;31m \u001b[0m\u001b[0;31m\\\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    242\u001b[0m             \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mstacked_obs\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mshape_dim0\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mclone\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~/Courses/UCSD/CSE257/pytorch-a2c-ppo-acktr-gail/a2c_ppo_acktr/envs.py\u001b[0m in \u001b[0;36mstep_wait\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m    185\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    186\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0mstep_wait\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\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--> 187\u001b[0;31m         \u001b[0mobs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mreward\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mdone\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0minfo\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mvenv\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mstep_wait\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    188\u001b[0m         \u001b[0mobs\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mfrom_numpy\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mobs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mfloat\u001b[0m\u001b[0;34m(\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[1;32m    189\u001b[0m         \u001b[0mreward\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mfrom_numpy\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mreward\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0munsqueeze\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mdim\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;36m1\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mfloat\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~/miniconda3/envs/ikos/lib/python3.8/site-packages/stable_baselines3/common/vec_env/dummy_vec_env.py\u001b[0m in \u001b[0;36mstep_wait\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m     49\u001b[0m                 \u001b[0mobs\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0menvs\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0menv_idx\u001b[0m\u001b[0;34m]\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     50\u001b[0m             \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_save_obs\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0menv_idx\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mobs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 51\u001b[0;31m         \u001b[0;32mreturn\u001b[0m \u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_obs_from_buf\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mnp\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcopy\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mbuf_rews\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mnp\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcopy\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mbuf_dones\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mdeepcopy\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mbuf_infos\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     52\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     53\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0mseed\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mseed\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mOptional\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mint\u001b[0m\u001b[0;34m]\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;34m->\u001b[0m \u001b[0mList\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mUnion\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;32mNone\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mint\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~/miniconda3/envs/ikos/lib/python3.8/copy.py\u001b[0m in \u001b[0;36mdeepcopy\u001b[0;34m(x, memo, _nil)\u001b[0m\n\u001b[1;32m    126\u001b[0m \u001b[0;32mdel\u001b[0m \u001b[0md\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mt\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    127\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 128\u001b[0;31m \u001b[0;32mdef\u001b[0m \u001b[0mdeepcopy\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mx\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mmemo\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;32mNone\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0m_nil\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[0m\u001b[1;32m    129\u001b[0m     \"\"\"Deep copy operation on arbitrary Python objects.\n\u001b[1;32m    130\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m: "
     ]
    }
   ],
   "source": [
    "for i, (prev, action, cur) in zip(range(20000), dataloader):\n",
    "    prev = prev.float().to(device)\n",
    "    action = action.float().to(device)\n",
    "    cur = cur.float().to(device)\n",
    "    pred = net(prev, action)\n",
    "    \n",
    "    optimizer.zero_grad()\n",
    "    out = RAMLoss(pred, cur)\n",
    "    out.backward()\n",
    "    optimizer.step()\n",
    "#     scheduler.step(out)\n",
    "    \n",
    "    if i % 100 == 0:\n",
    "        print(f'{i}, {out.item()}')\n",
    "    if out.item() < best_loss:\n",
    "        best_loss = out.item()\n",
    "        torch.save([net, best_loss], 'trained_models/cp/predictor.pt')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 134,
   "id": "employed-planning",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0, 84.97222900390625\n",
      "100, 244.71353149414062\n",
      "200, 122.55136108398438\n",
      "300, 77.8275146484375\n",
      "400, 61.221317291259766\n",
      "500, 82.46733093261719\n",
      "600, 113.71481323242188\n",
      "700, 103.4535903930664\n",
      "800, 92.49735260009766\n",
      "900, 100.73908996582031\n",
      "1000, 114.33695220947266\n",
      "1100, 98.61327362060547\n",
      "1200, 98.59130096435547\n",
      "1300, 89.2513427734375\n",
      "1400, 104.36087036132812\n",
      "1500, 89.0654296875\n",
      "1600, 100.53943634033203\n",
      "1700, 133.02536010742188\n",
      "1800, 121.65409088134766\n",
      "1900, 90.86433410644531\n",
      "2000, 98.9177474975586\n",
      "2100, 135.02392578125\n",
      "2200, 150.87948608398438\n",
      "2300, 100.12068176269531\n",
      "2400, 109.56149291992188\n",
      "2500, 108.75056457519531\n",
      "2600, 94.17854309082031\n",
      "2700, 110.979248046875\n",
      "2800, 127.96937561035156\n",
      "2900, 106.25205993652344\n",
      "3000, 106.06056213378906\n",
      "3100, 104.81507873535156\n",
      "3200, 103.19715881347656\n",
      "3300, 95.39431762695312\n",
      "3400, 205.64013671875\n",
      "3500, 91.30162048339844\n",
      "3600, 97.04886627197266\n",
      "3700, 109.35940551757812\n",
      "3800, 101.76792907714844\n",
      "3900, 262.92572021484375\n",
      "4000, 98.69830322265625\n",
      "4100, 95.42799377441406\n",
      "4200, 98.83941650390625\n"
     ]
    },
    {
     "ename": "KeyboardInterrupt",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-134-f962f2f68857>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[0;32m----> 1\u001b[0;31m \u001b[0;32mfor\u001b[0m \u001b[0mi\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m(\u001b[0m\u001b[0mprev\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0maction\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mcur\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mzip\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mrange\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;36m20000\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mdataloader\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      2\u001b[0m     \u001b[0mprev\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mprev\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mfloat\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mto\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[1;32m      3\u001b[0m     \u001b[0maction\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0maction\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mfloat\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mto\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[1;32m      4\u001b[0m     \u001b[0mcur\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mcur\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mfloat\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mto\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[1;32m      5\u001b[0m     \u001b[0mpred\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mnet\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mprev\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0maction\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/miniconda3/envs/ikos/lib/python3.8/site-packages/torch/utils/data/dataloader.py\u001b[0m in \u001b[0;36m__next__\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m    433\u001b[0m         \u001b[0;32mif\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_sampler_iter\u001b[0m \u001b[0;32mis\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    434\u001b[0m             \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_reset\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--> 435\u001b[0;31m         \u001b[0mdata\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_next_data\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    436\u001b[0m         \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_num_yielded\u001b[0m \u001b[0;34m+=\u001b[0m \u001b[0;36m1\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    437\u001b[0m         \u001b[0;32mif\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_dataset_kind\u001b[0m \u001b[0;34m==\u001b[0m \u001b[0m_DatasetKind\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mIterable\u001b[0m \u001b[0;32mand\u001b[0m\u001b[0;31m \u001b[0m\u001b[0;31m\\\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/miniconda3/envs/ikos/lib/python3.8/site-packages/torch/utils/data/dataloader.py\u001b[0m in \u001b[0;36m_next_data\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m    473\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0m_next_data\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\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    474\u001b[0m         \u001b[0mindex\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_next_index\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m  \u001b[0;31m# may raise StopIteration\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 475\u001b[0;31m         \u001b[0mdata\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_dataset_fetcher\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mfetch\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mindex\u001b[0m\u001b[0;34m)\u001b[0m  \u001b[0;31m# may raise StopIteration\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    476\u001b[0m         \u001b[0;32mif\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_pin_memory\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    477\u001b[0m             \u001b[0mdata\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0m_utils\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mpin_memory\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mpin_memory\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mdata\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/miniconda3/envs/ikos/lib/python3.8/site-packages/torch/utils/data/_utils/fetch.py\u001b[0m in \u001b[0;36mfetch\u001b[0;34m(self, possibly_batched_index)\u001b[0m\n\u001b[1;32m     26\u001b[0m             \u001b[0;32mfor\u001b[0m \u001b[0m_\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mpossibly_batched_index\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     27\u001b[0m                 \u001b[0;32mtry\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 28\u001b[0;31m                     \u001b[0mdata\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mappend\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mnext\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdataset_iter\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     29\u001b[0m                 \u001b[0;32mexcept\u001b[0m \u001b[0mStopIteration\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     30\u001b[0m                     \u001b[0;32mbreak\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m<ipython-input-103-1c49294b87d6>\u001b[0m in \u001b[0;36mexperience_gen\u001b[0;34m(env, agent)\u001b[0m\n\u001b[1;32m     13\u001b[0m             \u001b[0maction\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mrandint\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;36m0\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;36m6\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m(\u001b[0m\u001b[0;36m1\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0munsqueeze\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m-\u001b[0m\u001b[0;36m1\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     14\u001b[0m         \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 15\u001b[0;31m             \u001b[0m_\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0maction\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0m_\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0m_\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0magent\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mact\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mvis_obs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;32mNone\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     16\u001b[0m             \u001b[0maction\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0maction\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcpu\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     17\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/Courses/UCSD/CSE257/pytorch-a2c-ppo-acktr-gail/a2c_ppo_acktr/model.py\u001b[0m in \u001b[0;36mact\u001b[0;34m(self, inputs, rnn_hxs, masks, deterministic)\u001b[0m\n\u001b[1;32m     53\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     54\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0mact\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0minputs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mrnn_hxs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mmasks\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mdeterministic\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;32mFalse\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---> 55\u001b[0;31m         \u001b[0mvalue\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mactor_features\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mrnn_hxs\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mbase\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0minputs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mrnn_hxs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mmasks\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     56\u001b[0m         \u001b[0mdist\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdist\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mactor_features\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     57\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/miniconda3/envs/ikos/lib/python3.8/site-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    725\u001b[0m             \u001b[0mresult\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_slow_forward\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[1;32m    726\u001b[0m         \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 727\u001b[0;31m             \u001b[0mresult\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mforward\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    728\u001b[0m         for hook in itertools.chain(\n\u001b[1;32m    729\u001b[0m                 \u001b[0m_global_forward_hooks\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mvalues\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~/Courses/UCSD/CSE257/pytorch-a2c-ppo-acktr-gail/a2c_ppo_acktr/model.py\u001b[0m in \u001b[0;36mforward\u001b[0;34m(self, inputs, rnn_hxs, masks)\u001b[0m\n\u001b[1;32m    188\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    189\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[0minputs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mrnn_hxs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mmasks\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--> 190\u001b[0;31m         \u001b[0mx\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mmain\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0minputs\u001b[0m \u001b[0;34m/\u001b[0m \u001b[0;36m255.0\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    191\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    192\u001b[0m         \u001b[0;32mif\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mis_recurrent\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/miniconda3/envs/ikos/lib/python3.8/site-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    725\u001b[0m             \u001b[0mresult\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_slow_forward\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[1;32m    726\u001b[0m         \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 727\u001b[0;31m             \u001b[0mresult\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mforward\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    728\u001b[0m         for hook in itertools.chain(\n\u001b[1;32m    729\u001b[0m                 \u001b[0m_global_forward_hooks\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mvalues\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~/miniconda3/envs/ikos/lib/python3.8/site-packages/torch/nn/modules/container.py\u001b[0m in \u001b[0;36mforward\u001b[0;34m(self, input)\u001b[0m\n\u001b[1;32m    115\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[0minput\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[0;32mfor\u001b[0m \u001b[0mmodule\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mself\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[0minput\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mmodule\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0minput\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;32mreturn\u001b[0m \u001b[0minput\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    119\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/miniconda3/envs/ikos/lib/python3.8/site-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    725\u001b[0m             \u001b[0mresult\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_slow_forward\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[1;32m    726\u001b[0m         \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 727\u001b[0;31m             \u001b[0mresult\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mforward\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    728\u001b[0m         for hook in itertools.chain(\n\u001b[1;32m    729\u001b[0m                 \u001b[0m_global_forward_hooks\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mvalues\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~/miniconda3/envs/ikos/lib/python3.8/site-packages/torch/nn/modules/activation.py\u001b[0m in \u001b[0;36mforward\u001b[0;34m(self, input)\u001b[0m\n\u001b[1;32m    100\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    101\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[0minput\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mTensor\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;34m->\u001b[0m \u001b[0mTensor\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[0;32mreturn\u001b[0m \u001b[0mF\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mrelu\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0minput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0minplace\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0minplace\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;32mdef\u001b[0m \u001b[0mextra_repr\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;34m->\u001b[0m \u001b[0mstr\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/miniconda3/envs/ikos/lib/python3.8/site-packages/torch/nn/functional.py\u001b[0m in \u001b[0;36mrelu\u001b[0;34m(input, inplace)\u001b[0m\n\u001b[1;32m   1134\u001b[0m         \u001b[0mresult\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mrelu_\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0minput\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1135\u001b[0m     \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 1136\u001b[0;31m         \u001b[0mresult\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mrelu\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0minput\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   1137\u001b[0m     \u001b[0;32mreturn\u001b[0m \u001b[0mresult\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1138\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m: "
     ]
    }
   ],
   "source": [
    "for i, (prev, action, cur) in zip(range(20000), dataloader):\n",
    "    prev = prev.float().to(device)\n",
    "    action = action.float().to(device)\n",
    "    cur = cur.float().to(device)\n",
    "    pred = net(prev, action)\n",
    "    \n",
    "    optimizer.zero_grad()\n",
    "    out = RAMLoss(pred, cur)\n",
    "    out.backward()\n",
    "    optimizer.step()\n",
    "#     scheduler.step(out)\n",
    "    \n",
    "    if i % 100 == 0:\n",
    "        print(f'{i}, {out.item()}')\n",
    "    if out.item() < best_loss:\n",
    "        best_loss = out.item()\n",
    "        torch.save([net, best_loss], 'trained_models/cp/predictor.pt')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 69,
   "id": "confident-border",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([186.3287, 186.1859, 185.9084, 185.6812, 185.3882, 185.1511, 184.8581,\n",
       "        184.6210, 184.3173, 184.0700, 183.7642, 183.5169, 183.2110, 174.4593,\n",
       "        174.8849, 171.5093, 171.7436, 175.7138, 175.9613, 172.4207, 172.6434,\n",
       "        176.7037, 176.9512, 173.2953, 173.5764, 177.6667, 174.1609, 174.3833,\n",
       "        178.4554, 178.7112, 175.0382, 175.2431, 182.2193, 186.8036, 186.6889,\n",
       "        186.6277, 186.4290, 186.2093, 185.9328, 185.7132, 185.4308, 185.2050,\n",
       "        184.9176, 184.6914, 184.3986, 184.1566, 183.8544, 183.6124, 174.6291,\n",
       "        175.0713, 171.7400, 172.0045, 175.9480, 176.2103, 172.7980, 173.0378,\n",
       "        176.9944, 177.2484, 173.7389, 174.0315, 178.0106, 174.3979, 174.5976,\n",
       "        178.7258], device='cuda:0', grad_fn=<SelectBackward>)"
      ]
     },
     "execution_count": 69,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "pred[:, 49]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "id": "quick-wilson",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_ball_pos(ram):\n",
    "    ball_x = ram[49]  # X coordinate of ball\n",
    "    ball_y = ram[54]\n",
    "    return ball_x, ball_y"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 56,
   "id": "amazing-project",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(tensor(82, device='cuda:0', dtype=torch.int32), tensor(44, device='cuda:0', dtype=torch.int32))\n",
      "(tensor(52., device='cuda:0'), tensor(0., device='cuda:0'))\n"
     ]
    }
   ],
   "source": [
    "print(get_ball_pos(pred[13].int()))\n",
    "print(get_ball_pos(cur[13]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "compressed-torture",
   "metadata": {},
   "outputs": [],
   "source": [
    "# ram = env.unwrapped._get_ram()  # get emulator RAM state\n",
    "# cpu_score = ram[13]  # computer/ai opponent score \n",
    "# player_score = ram[14]  # your score\n",
    "# cpu_paddle_y = ram[21]  # Y coordinate of computer paddle\n",
    "# player_paddle_y = ram[51]  # Y coordinate of your paddle\n",
    "# ball_x = ram[49]  # X coordinate of ball\n",
    "# ball_y = ram[54]  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "warming-lyric",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}