frame_stacked_cartpole.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "id": "h6v4c9_UwOmr",
    "outputId": "d5b5140a-cd61-45f6-859d-79986d112476"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The tensorboard extension is already loaded. To reload it, use:\n",
      "  %reload_ext tensorboard\n"
     ]
    }
   ],
   "source": [
    "%load_ext tensorboard\n",
    "\n",
    "import math\n",
    "import gym \n",
    "import plotly.express as px\n",
    "import numpy as np\n",
    "import warnings\n",
    "from plotly.subplots import make_subplots\n",
    "import plotly.graph_objects as go\n",
    "import random\n",
    "from collections import Counter, deque\n",
    "import pandas as pd\n",
    "from sklearn.model_selection import train_test_split\n",
    "import datetime\n",
    "from sklearn.metrics import confusion_matrix\n",
    "import plotly.figure_factory as ff\n",
    "\n",
    "import os\n",
    "import glob\n",
    "from IPython import display as ipythondisplay\n",
    "from tqdm.notebook import tqdm\n",
    "from gym.wrappers import Monitor\n",
    "from IPython.display import HTML\n",
    "import base64\n",
    "import io\n",
    "import pickle\n",
    "import torch\n",
    "from torch import nn\n",
    "import kornia.augmentation as kaug\n",
    "warnings.filterwarnings(\"ignore\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "device(type='cuda')"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
    "device"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "id": "XREqct6awOms"
   },
   "outputs": [],
   "source": [
    "from pyvirtualdisplay import Display\n",
    "display = Display(visible=False, size=(1400, 900))\n",
    "if torch.cuda.is_available():\n",
    "    display.start()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "id": "T1ACavBgwOmt"
   },
   "outputs": [],
   "source": [
    "import os, sys, copy, argparse, shutil\n",
    "def parse_arguments():\n",
    "    parser = argparse.ArgumentParser(description='Deep Q Network Argument Parser')\n",
    "    parser.add_argument('--seed', dest='seed', type=int, default=1)\n",
    "    parser.add_argument('--env', dest='env', type=str, default='CartPole-v0')\n",
    "    parser.add_argument('--save_interval', type=int, default=50, help='save model every n episodes')\n",
    "    parser.add_argument('--log_interval', type=int, default=10, help='logging every n episodes')\n",
    "    parser.add_argument('--render', help='render', type=int, default=1)\n",
    "    parser.add_argument('--batch_size', help='batch_size', type=int, default=32)\n",
    "    parser.add_argument('--train_freq', help='train_frequency', type=int, default=1)\n",
    "    parser.add_argument('--max_episode', help='maximum episode', type=int, default=None)\n",
    "    parser.add_argument('--max_timesteps', help='maximum timestep', type=int, default=100000000)\n",
    "    parser.add_argument('--lr', dest='lr', type=float, default=0.00025)\n",
    "    parser.add_argument('--lr_decay', action='store_true', help='decay learning rate')\n",
    "    parser.add_argument('--gamma', help='discount_factor', type=float, default=0.99)\n",
    "    parser.add_argument('--warmup_mem', type=int, help='warmup memory size', default=1000)\n",
    "    parser.add_argument('--frame_skip', type=int, help='number of frames to skip for each action', default=3)\n",
    "    parser.add_argument('--frame_stack', type=int, help='number of frames to stack', default=4)\n",
    "    parser.add_argument('--memory', help='memory size', type=int, default=1000000)\n",
    "    parser.add_argument('--initial_epsilon', '-ie', help='initial_epsilon', type=float, default=0.5)\n",
    "    parser.add_argument('--final_epsilon', '-fe', help='final_epsilon', type=float, default=0.05)\n",
    "    parser.add_argument('--max_epsilon_decay_steps', '-eds', help='maximum steps to decay epsilon', type=int, default=100000)\n",
    "    parser.add_argument('--max_grad_norm', type=float, default=None, help='maximum gradient norm')\n",
    "    parser.add_argument('--soft_update', '-su', action='store_true', help='soft update target network')\n",
    "    parser.add_argument('--double_q', '-dq', action='store_true', help='enabling double DQN')\n",
    "    parser.add_argument('--dueling_net', '-dn', action='store_true', help='enabling dueling network')\n",
    "    parser.add_argument('--test', action='store_true', help='test the trained model')\n",
    "    parser.add_argument('--tau', type=float, default=0.01, help='tau for soft target network update')\n",
    "    parser.add_argument('--hard_update_freq', '-huf', type=int, default=2000, help='hard target network update frequency')\n",
    "    parser.add_argument('--save_dir', type=str, default='./data')\n",
    "    parser.add_argument('--resume_step', '-rs', type=int, default=None)\n",
    "    return parser.parse_args()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "id": "X3Rkxe4VwOmu"
   },
   "outputs": [],
   "source": [
    "#@title Set up constants for env and training\n",
    "test = False \n",
    "save_dir = './data'\n",
    "render = False\n",
    "max_episode = None\n",
    "max_timesteps = 100000000\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "id": "539HMtWFwOmv"
   },
   "outputs": [
    {
     "ename": "RuntimeError",
     "evalue": "CUDA error: out of memory\nCUDA kernel errors might be asynchronously reported at some other API call,so the stacktrace below might be incorrect.\nFor debugging consider passing CUDA_LAUNCH_BLOCKING=1.",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mRuntimeError\u001b[0m                              Traceback (most recent call last)",
      "\u001b[0;32m/tmp/ipykernel_4587/2163621172.py\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m     21\u001b[0m \u001b[0mrandom_elastic_transform\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mkaug\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mRandomElasticTransform\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     22\u001b[0m random_fisheye = kaug.RandomFisheye(\n\u001b[0;32m---> 23\u001b[0;31m         \u001b[0mcenter_x\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mtorch\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[0;36m.3\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;36m.3\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[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     24\u001b[0m         \u001b[0mcenter_y\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mtorch\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[0;36m.3\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;36m.3\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[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     25\u001b[0m         \u001b[0mgamma\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtensor\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;36m.9\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[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[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mRuntimeError\u001b[0m: CUDA error: out of memory\nCUDA kernel errors might be asynchronously reported at some other API call,so the stacktrace below might be incorrect.\nFor debugging consider passing CUDA_LAUNCH_BLOCKING=1."
     ]
    }
   ],
   "source": [
    "#@title Augmentations\n",
    "'''\n",
    "color_jitter\n",
    "random_elastic_transform\n",
    "random_fisheye\n",
    "random_color_equalize\n",
    "random_gaussian_blur\n",
    "random_gaussian_noise\n",
    "random_horizontal_flip\n",
    "random_color_invert\n",
    "random_perspective_shift\n",
    "random_shift\n",
    "'''\n",
    "color_jitter = kaug.ColorJitter(\n",
    "        brightness=np.random.random(),\n",
    "        contrast=np.random.random(),\n",
    "        saturation=np.random.random(),\n",
    "        hue=np.random.random(),\n",
    "        p=0.5\n",
    "        )\n",
    "random_elastic_transform = kaug.RandomElasticTransform()\n",
    "random_fisheye = kaug.RandomFisheye(\n",
    "        center_x=torch.tensor([-.3, .3]).to(device),\n",
    "        center_y=torch.tensor([-.3, .3]).to(device),\n",
    "        gamma=torch.tensor([.9, 1.]).to(device),\n",
    "        )\n",
    "# need to divide by 255.0\n",
    "random_color_equalize = lambda obs: kaug.RandomEqualize()(obs / 255.) * 255\n",
    "random_gaussian_blur = kaug.RandomGaussianBlur(\n",
    "        kernel_size=(9, 9),\n",
    "        sigma = (5., 5.)\n",
    "        )\n",
    "random_gaussian_noise = kaug.RandomGaussianNoise()\n",
    "random_horizontal_flip = kaug.RandomHorizontalFlip()\n",
    "random_color_invert = kaug.RandomInvert()\n",
    "random_perspective_shift = kaug.RandomPerspective()\n",
    "get_random_shift = lambda h, w, shift_by: nn.Sequential(kaug.RandomCrop((h - shift_by, w - shift_by)), nn.ReplicationPad2d(20), kaug.RandomCrop((h - shift_by, w - shift_by)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "K78NVCz-oqHc"
   },
   "outputs": [],
   "source": [
    "def tie_weights(src, trg):\n",
    "    assert type(src) == type(trg)\n",
    "    trg.weight = src.weight\n",
    "    trg.bias = src.bias\n",
    "\n",
    "def conv_output_shape(h_w, kernel_size=1, stride=1, pad=0, dilation=1):\n",
    "    from math import floor\n",
    "    if type(kernel_size) is not tuple:\n",
    "        kernel_size = (kernel_size, kernel_size)\n",
    "    h = floor( ((h_w[0] + (2 * pad) - ( dilation * (kernel_size[0] - 1) ) - 1 )/ stride) + 1)\n",
    "    w = floor( ((h_w[1] + (2 * pad) - ( dilation * (kernel_size[1] - 1) ) - 1 )/ stride) + 1)\n",
    "    return h, w\n",
    "\n",
    "# for 84 x 84 inputs\n",
    "OUT_DIM = {2: 39, 4: 35, 6: 31}\n",
    "# for 64 x 64 inputs\n",
    "OUT_DIM_64 = {2: 29, 4: 25, 6: 21}\n",
    "\n",
    "''' TODO change the layer parameters ''' \n",
    "class PixelEncoder(nn.Module):\n",
    "    \"\"\"Convolutional encoder of pixels observations.\"\"\"\n",
    "    def __init__(self, obs_shape, input_channels, feature_dim=50, num_layers=3, num_filters=64, output_logits=False):\n",
    "        super().__init__()\n",
    "\n",
    "        assert len(obs_shape) == 3\n",
    "        self.obs_shape = (obs_shape[2], obs_shape[0], obs_shape[1])\n",
    "        self.feature_dim = feature_dim\n",
    "        self.num_layers = num_layers\n",
    "        \n",
    "        # 160, 210, 3\n",
    "        self.conv1 = nn.Conv2d(input_channels, 32, kernel_size=5, stride=5)\n",
    "        conv1_shape = conv_output_shape(self.obs_shape[1:], kernel_size=5, stride=5)\n",
    "        # Input to conv2: 32, 42, 32\n",
    "        self.conv2 = nn.Conv2d(32, 64, kernel_size=4, stride=2)\n",
    "        conv2_shape = conv_output_shape(conv1_shape, kernel_size=4, stride=2)\n",
    "        # Input to conv3: 15, 20, 64\n",
    "        self.conv3 = nn.Conv2d(64, 64, kernel_size=4, stride=1)\n",
    "        conv3_shape = conv_output_shape(conv2_shape, kernel_size=4, stride=1)\n",
    "        # Output from conv3: 12, 17, 64\n",
    "\n",
    "        # out_dim = OUT_DIM_64[num_layers] if obs_shape[-1] == 64 else OUT_DIM[num_layers]\n",
    "        out_dims = conv3_shape\n",
    "        self.fc = nn.Linear(num_filters * out_dims[0] * out_dims[1], self.feature_dim)\n",
    "        self.ln = nn.LayerNorm(self.feature_dim)\n",
    "\n",
    "        self.outputs = dict()\n",
    "        self.output_logits = output_logits\n",
    "\n",
    "    def reparameterize(self, mu, logstd):\n",
    "        std = torch.exp(logstd)\n",
    "        eps = torch.randn_like(std)\n",
    "        return mu + eps * std\n",
    "\n",
    "    def forward_conv(self, obs):\n",
    "        self.outputs['obs'] = obs\n",
    "        conv1 = torch.relu(self.conv1(obs))\n",
    "        self.outputs['conv1'] = conv1\n",
    "        conv2 = torch.relu(self.conv2(conv1))\n",
    "        self.outputs['conv2'] = conv2\n",
    "        conv3 = torch.relu(self.conv3(conv2))\n",
    "        self.outputs['conv3'] = conv3\n",
    "\n",
    "        h = conv3.reshape(conv3.size(0), -1)\n",
    "        return h\n",
    "\n",
    "    def forward(self, obs, detach=False):\n",
    "        h = self.forward_conv(obs)\n",
    "\n",
    "        if detach:\n",
    "            h = h.detach()\n",
    "\n",
    "        h_fc = self.fc(h)\n",
    "        self.outputs['fc'] = h_fc\n",
    "\n",
    "        h_norm = self.ln(h_fc)\n",
    "        self.outputs['ln'] = h_norm\n",
    "\n",
    "        if self.output_logits:\n",
    "            out = h_norm\n",
    "        else:\n",
    "            out = torch.tanh(h_norm)\n",
    "            self.outputs['tanh'] = out\n",
    "\n",
    "        return out\n",
    "\n",
    "    def copy_conv_weights_from(self, source):\n",
    "        \"\"\"Tie convolutional layers\"\"\"\n",
    "        # only tie conv layers\n",
    "        for i in range(self.num_layers):\n",
    "            tie_weights(src=source.convs[i], trg=self.convs[i])\n",
    "\n",
    "    def log(self, L, step, log_freq):\n",
    "        if step % log_freq != 0:\n",
    "            return\n",
    "\n",
    "        for k, v in self.outputs.items():\n",
    "            L.log_histogram('train_encoder/%s_hist' % k, v, step)\n",
    "            if len(v.shape) > 2:\n",
    "                L.log_image('train_encoder/%s_img' % k, v[0], step)\n",
    "\n",
    "        for i in range(self.num_layers):\n",
    "            L.log_param('train_encoder/conv%s' % (i + 1), self.convs[i], step)\n",
    "        L.log_param('train_encoder/fc', self.fc, step)\n",
    "        L.log_param('train_encoder/ln', self.ln, step)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "0kIzyV1QwOmw"
   },
   "outputs": [],
   "source": [
    "def wrap_env(env, train=True):\n",
    "    suffix = 'train' if train else 'test'\n",
    "    monitor_dir = os.path.join(save_dir, 'monitor_%s' % suffix)\n",
    "    os.makedirs(monitor_dir, exist_ok=True)\n",
    "    if not train:\n",
    "        video_save_interval = 10\n",
    "        env = Monitor(env, directory=monitor_dir,\n",
    "                      video_callable=lambda episode_id: episode_id % video_save_interval == 0,\n",
    "                      force=True)\n",
    "    else:\n",
    "        if render:\n",
    "            if max_episode is not None:\n",
    "                video_save_interval = int(max_episode / 3)\n",
    "            else:\n",
    "                video_save_interval = int(max_timesteps / float(env._max_episode_steps) / 3)\n",
    "            env = Monitor(env, directory=monitor_dir,\n",
    "                          video_callable=lambda episode_id: episode_id % video_save_interval == 0,\n",
    "                          force=True)\n",
    "        else:\n",
    "            env = Monitor(env, directory=monitor_dir, video_callable=False, force=True)\n",
    "    return env"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "mZgEGp83wOmw"
   },
   "outputs": [],
   "source": [
    "class ReplayMemory(object):\n",
    "    def __init__(self, max_epi_num=100, max_epi_len=200, obs_shape=(210, 160)):\n",
    "        # capacity is the maximum number of steps in memory\n",
    "        self.max_epi_num = max_epi_num\n",
    "        self.max_epi_len = max_epi_len\n",
    "        # saves each tuple of (state, action, next state, reward)\n",
    "        self.capacity = 100 # self.max_epi_num * max_epi_len\n",
    "        self.idx = 0\n",
    "        # Use 6 frames (stacking 2 x 3 frame obs)\n",
    "        self.obs_memory = np.zeros((self.capacity, *obs_shape, 6)) # deque(maxlen=self.max_epi_num * max_epi_len)\n",
    "        self.next_memory = np.zeros((self.capacity, *obs_shape, 6))\n",
    "        self.act_memory = np.zeros((self.capacity, 1))\n",
    "        self.reward_memory = np.zeros((self.capacity, 1))\n",
    "        self.is_av = False\n",
    "        self.current_epi = 0\n",
    "\n",
    "    def reset(self):\n",
    "        self.current_epi = 0\n",
    "        self.memory.clear()\n",
    "\n",
    "    ''' deprecated for tuple buffer '''\n",
    "    def create_new_epi(self):\n",
    "        pass\n",
    "\n",
    "    def remember(self, state, next_state, action, reward):\n",
    "        idx = self.idx % self.capacity\n",
    "        self.obs_memory[idx] = state.copy()\n",
    "        self.next_memory[idx] = next_state.copy()\n",
    "        self.act_memory[idx] = action\n",
    "        self.reward_memory[idx] = reward\n",
    "        self.idx += 1\n",
    "\n",
    "                \n",
    "    # samples batch_size\n",
    "    def sample(self, batch_size):\n",
    "        if batch_size < self.idx:\n",
    "            max_len = min(self.capacity, self.idx)\n",
    "            idx = np.random.randint(0, max_len - 1, batch_size)\n",
    "            return self.obs_memory[idx], self.next_memory[idx], self.act_memory[idx], self.reward_memory[idx]\n",
    "        return self.obs_memory[:self.idx], self.next_memory[:self.idx], self.act_memory[:self.idx], self.reward_memory[:self.idx]\n",
    "\n",
    "    def size(self):\n",
    "        return self.idx\n",
    "\n",
    "    def is_available(self):\n",
    "        self.is_av = True\n",
    "        if self.idx <= 1:\n",
    "            self.is_av = False\n",
    "        return self.is_av\n",
    "\n",
    "    def print_info(self):\n",
    "        pass"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "FmR5zvSVwOmw"
   },
   "outputs": [],
   "source": [
    "#@title Create a training conv agent\n",
    "import torch.nn.functional as F\n",
    "\n",
    "class DQNetworkConv(nn.Module):\n",
    "    def __init__(self, in_channels, act_dim, dueling=False):\n",
    "        super(DQNetworkConv, self).__init__()\n",
    "        self.act_dim = act_dim\n",
    "        self.dueling = dueling\n",
    "        # 160, 210, 3 \n",
    "        self.conv1 = nn.Conv2d(in_channels, 32, kernel_size=5, stride=5)\n",
    "        # Input to conv2: 32, 42, 32\n",
    "        self.conv2 = nn.Conv2d(32, 64, kernel_size=4, stride=2)\n",
    "        # Input to conv3: 15, 20, 64\n",
    "        self.conv3 = nn.Conv2d(64, 64, kernel_size=4, stride=1)\n",
    "        # Output from conv3: 12, 17, 64\n",
    "        if self.dueling:\n",
    "            self.v_fc4 = nn.Linear(12 * 17 * 64, 512)\n",
    "            self.adv_fc4 = nn.Linear(12 * 17 * 64, 512)\n",
    "            self.v_fc5 = nn.Linear(512, 1)\n",
    "            self.adv_fc5 = nn.Linear(512, self.act_dim)\n",
    "        else:\n",
    "            self.fc4 = nn.Linear(12 * 17 * 64, 512)\n",
    "            self.fc5 = nn.Linear(512, self.act_dim)\n",
    "        self.parameters = (list(self.conv1.parameters())) + (list(self.conv2.parameters())) + (list(self.conv3.parameters())) + (list(self.fc4.parameters())) + (list(self.fc5.parameters()))\n",
    "\n",
    "    def forward(self, st):\n",
    "        out = F.relu(self.conv1(st))\n",
    "        out = F.relu(self.conv2(out))\n",
    "        out = F.relu(self.conv3(out))\n",
    "        out = out.view(out.size(0), -1)\n",
    "        if self.dueling:\n",
    "            val = F.relu(self.v_fc4(out))\n",
    "            adv = F.relu(self.adv_fc4(out))\n",
    "            val = self.v_fc5(val)\n",
    "            adv = self.adv_fc5(adv)\n",
    "            out = val.expand_as(adv) + adv - adv.mean(-1, keepdim=True).expand_as(adv)\n",
    "        else:\n",
    "            out = F.relu(self.fc4(out))\n",
    "            out = self.fc5(out)\n",
    "        return out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "ZGsarcvPzoNm"
   },
   "outputs": [],
   "source": [
    "#@title Create a training FC agent\n",
    "import torch.nn.functional as F\n",
    "\n",
    "class DQNetworkFC(nn.Module):\n",
    "    def __init__(self, z_dim, act_dim, dueling=False):\n",
    "        super(DQNetworkFC, self).__init__()\n",
    "        self.act_dim = act_dim\n",
    "        self.input_dim = z_dim \n",
    "        self.dueling = dueling\n",
    "        if self.dueling:\n",
    "            self.v_fc1 = nn.Linear(z_dim, 512)\n",
    "            self.adv_fc1 = nn.Linear(z_dim, 512)\n",
    "            self.v_fc2 = nn.Linear(512, 1)\n",
    "            self.adv_fc2 = nn.Linear(512, 256)\n",
    "            self.v_fc3 = nn.Linear(256, 1)\n",
    "            self.adv_fc3 = nn.Linear(256, self.act_dim)\n",
    "        else:\n",
    "            self.fc1 = nn.Linear(z_dim, 512)\n",
    "            self.fc2 = nn.Linear(512, 256)\n",
    "            self.fc3 = nn.Linear(256, self.act_dim)\n",
    "\n",
    "    def forward(self, st):\n",
    "        out = F.relu(self.fc1(st))\n",
    "        out = F.relu(self.fc2(out))\n",
    "        ''' Do we need a relu on the last layer if the output is probability over action space? '''\n",
    "        out = F.relu(self.fc3(out))\n",
    "        return out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "I5dCMf6CH99A"
   },
   "outputs": [],
   "source": [
    "def process_obs(obs, divide=False):\n",
    "    obs = torch.Tensor(obs / 255. if divide else obs)\n",
    "    if len(obs.shape) < 4:\n",
    "        obs = obs.unsqueeze(0)\n",
    "    obs = obs.permute(0, 3, 1, 2)\n",
    "    return obs.to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "H6_wThKowOmx"
   },
   "outputs": [],
   "source": [
    "def take_action(env, action):\n",
    "    state, rew, done, _ = env.step(action)\n",
    "    obs = env.render(mode='rgb_array')\n",
    "    return obs, rew, done"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "XL9D4MCxwOmx"
   },
   "outputs": [],
   "source": [
    "MAX_STEPS = 200"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "9NKDLpKGJ2lB"
   },
   "outputs": [],
   "source": [
    "class CURL(nn.Module):\n",
    "    \"\"\"\n",
    "    CURL\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, obs_shape, z_dim, batch_size, encoder, output_type=\"continuous\", critic=None, critic_target=None):\n",
    "        super(CURL, self).__init__()\n",
    "        self.obs_shape = obs_shape\n",
    "        self.batch_size = batch_size\n",
    "\n",
    "        # self.encoder = critic.encoder\n",
    "        self.encoder = encoder \n",
    "\n",
    "        # self.encoder_target = critic_target.encoder \n",
    "        self.fc1 = nn.Linear(100, 50)\n",
    "        self.fc2 = nn.Linear(50, 1)\n",
    "\n",
    "        self.W = nn.Parameter(torch.rand(z_dim, z_dim))\n",
    "        self.output_type = output_type\n",
    "\n",
    "    def encode(self, x, detach=False, ema=False):\n",
    "        \"\"\"\n",
    "        Encoder: z_t = e(x_t)\n",
    "        :param x: x_t, x y coordinates\n",
    "        :return: z_t, value in r2\n",
    "        \"\"\"\n",
    "        if ema:\n",
    "            with torch.no_grad():\n",
    "                z_out = self.encoder_target(x)\n",
    "        else:\n",
    "            z_out = self.encoder(x)\n",
    "\n",
    "        if detach:\n",
    "            z_out = z_out.detach()\n",
    "        return z_out\n",
    "\n",
    "    def compute_logits(self, z_a, z_mod):\n",
    "        \"\"\"\n",
    "        Uses logits trick for CURL:\n",
    "        - compute (B,B) matrix z_a (W z_pos.T)\n",
    "        - positives are all diagonal elements\n",
    "        - negatives are all other elements\n",
    "        - to compute loss use multiclass cross entropy with identity matrix for labels\n",
    "        \"\"\"\n",
    "#         Wz = torch.matmul(self.W, z_mod.T)  # (z_dim,B)\n",
    "#         logits = torch.matmul(z_a, Wz)  # (B,B)\n",
    "#         logits = logits - torch.max(logits, 1)[0][:, None]\n",
    "#         return logits\n",
    "        input_zs = torch.cat([z_a, z_mod], 1)\n",
    "        logits = F.relu(self.fc1(input_zs))\n",
    "        logits = F.sigmoid(self.fc2(logits))\n",
    "        return logits"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "1wbEY4hTshF6"
   },
   "outputs": [],
   "source": [
    "#@title Generate a batch of negatively labelled examples given observations\n",
    "\n",
    "def generate_negatives(obs):\n",
    "    neg_idx = np.random.randint(len(obs), size=len(obs))\n",
    "    pos_idx = np.arange(len(obs))\n",
    "    resample = (neg_idx == pos_idx)\n",
    "    for (i, r) in enumerate(resample):\n",
    "        if r:\n",
    "            idx = neg_idx[i]\n",
    "        else:\n",
    "            idx = np.random.randint(0, len(obs), 1)[0]\n",
    "            while idx == i:\n",
    "                idx = np.random.randint(0, len(obs), 1)[0]\n",
    "        neg_idx[i] = idx\n",
    "    return (obs[neg_idx]).copy()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "x3KEWmwUwOmx"
   },
   "outputs": [],
   "source": [
    "#@title Create a training agent (wrapper for conv agent)\n",
    "\n",
    "GAMMA = 0.99\n",
    "\n",
    "class Agent(object):\n",
    "    def __init__(self, act_dim, in_channels=6, max_epi_num=50, max_epi_len=300, CURL=None, aug=None, conv_net=False, random_shift=None):\n",
    "        self.N_action = act_dim\n",
    "        self.max_epi_num = max_epi_num\n",
    "        self.max_epi_len = max_epi_len\n",
    "        ''' To decide when to copy weights to the target network '''\n",
    "        self.num_param_updates = 0\n",
    "        self.CURL = CURL\n",
    "        self.aug = aug\n",
    "        self.random_shift = random_shift\n",
    "        if conv_net:\n",
    "            self.conv_net = DQNetworkConv(in_channels, act_dim).to(device)\n",
    "            self.target = DQNetworkConv(in_channels, act_dim).to(device)\n",
    "        else:\n",
    "            ''' if using the encoder head for contrastive loss '''\n",
    "            self.conv_net = DQNetworkFC(self.CURL.encoder.feature_dim, act_dim).to(device)\n",
    "            self.target = DQNetworkFC(self.CURL.encoder.feature_dim, act_dim).to(device)\n",
    "        self.buffer = ReplayMemory(max_epi_num=self.max_epi_num, max_epi_len=self.max_epi_len, obs_shape=CURL.obs_shape[:2])\n",
    "        self.gamma = 0.99\n",
    "        self.loss_fn = torch.nn.MSELoss()\n",
    "        self.optimizer = torch.optim.Adam(\n",
    "            list(self.conv_net.parameters()) + \n",
    "            list(self.CURL.parameters()) + \n",
    "            list(self.CURL.encoder.parameters()), lr=1e-3)\n",
    "\n",
    "    def remember(self, state, action, reward, next_state):\n",
    "        self.buffer.remember(state, next_state, action, reward)\n",
    "\n",
    "    ''' Copy the weights to the target network every 100 updates '''\n",
    "    def train(self, batch_size=32, target_update_freq=100, use_encoder=True):\n",
    "        if self.buffer.is_available():\n",
    "            obs, next_obs, action_list, reward_list = self.buffer.sample(batch_size)\n",
    "            \n",
    "            ''' Pass through the encoder to get encodings\n",
    "             If also training the contrastive loss\n",
    "             include that here! \n",
    "             1. data augmentation to create pos and negative pairs\n",
    "             2. encoder \n",
    "             3. update encoder loss function (using a separate optimizer) or add to the loss computed below\n",
    "            '''\n",
    "\n",
    "            losses = []\n",
    "\n",
    "            # check if obs is a numpy or a torch tensor\n",
    "            if use_encoder:\n",
    "                obs_anchor = self.random_shift(process_obs(obs.copy()))\n",
    "                obs_pos = self.random_shift(process_obs(obs.copy()))\n",
    "                mixed_obs = generate_negatives(obs)\n",
    "                mixed_obs = process_obs(mixed_obs)\n",
    "                obs_neg = self.random_shift(mixed_obs)\n",
    "                if self.aug is not None:\n",
    "                    obs_pos = self.aug(obs_pos)\n",
    "                    obs_neg = self.aug(obs_neg)\n",
    "                z_a = self.CURL.encode(obs_anchor)\n",
    "                z_pos = self.CURL.encode(obs_pos)\n",
    "                # Mix pairs to generate negative labels\n",
    "                z_neg = self.CURL.encode(obs_neg)\n",
    "                next_obs = self.random_shift(process_obs(next_obs.copy()))\n",
    "                z_next = self.CURL.encode(next_obs)\n",
    "\n",
    "                # logits = self.CURL.compute_logits(z_a, z_pos)\n",
    "                # labels = torch.arange(logits.shape[0]).long().to(device)\n",
    "                pos_logits = self.CURL.compute_logits(z_a, z_pos)\n",
    "                neg_logits = self.CURL.compute_logits(z_a, z_neg)\n",
    "                # [32, 32]\n",
    "                pos_labels = torch.ones((pos_logits.shape[0], 1))# .long()\n",
    "                neg_labels = torch.zeros((neg_logits.shape[0], 1))# .long() \n",
    "                # TODO: stack pos and neg logits and labels (double check dim)\n",
    "                logits = torch.cat([pos_logits, neg_logits], 0)\n",
    "                labels = torch.cat([pos_labels, neg_labels], 0).to(device)\n",
    "                \n",
    "                # pass into the loss function\n",
    "                # encoding_loss = nn.CrossEntropyLoss()(logits, labels)\n",
    "                encoding_loss = nn.BCELoss()(logits, labels)\n",
    "\n",
    "                ''' Combine encoding loss with rl loss below '''\n",
    "                losses.append(encoding_loss)\n",
    "\n",
    "                # Then pass that encoding through the conv_net to get Q value estimates\n",
    "                Qs = self.conv_net(z_a)\n",
    "                next_Qs = self.target(z_next).detach().max(1)[0]\n",
    "            \n",
    "            else:\n",
    "                ''' If not using the encoder, pass the obs directly to the CNN '''\n",
    "                obs = process_obs(obs)\n",
    "                # estimate current q values from observations\n",
    "                Qs = self.conv_net(obs)\n",
    "                # find next max q values based on next observations\n",
    "                next_Qs = self.target(next_obs).detach().max(1)[0]\n",
    "            \n",
    "            ''' find target q values ''' \n",
    "            next_Qs = next_Qs.cpu().numpy() \n",
    "            Qs = torch.gather(Qs, dim=1, index=torch.tensor(action_list, dtype=torch.int64).to(device)).float().to(device)\n",
    "            target_Qs = torch.tensor(reward_list.squeeze(-1) + GAMMA * next_Qs).float().to(device)\n",
    "            ''' try to set Qs equal to target_Qs '''\n",
    "            q_loss = self.loss_fn(Qs, target_Qs)\n",
    "            losses.append(q_loss)\n",
    "            \n",
    "            ''' Loss update for q network and encoder head '''\n",
    "            losses = torch.stack(losses).sum()\n",
    "            self.optimizer.zero_grad()\n",
    "            losses.backward()\n",
    "            self.optimizer.step()\n",
    "\n",
    "            self.num_param_updates += 1\n",
    "            if self.num_param_updates % target_update_freq == 0:\n",
    "                self.target.load_state_dict(self.conv_net.state_dict())\n",
    "\n",
    "    # TODO: check the sizes of inputs and outputs\n",
    "    def get_action(self, obs, epsilon, use_encoding=True):\n",
    "        ''' \n",
    "         If using an encoder, need to pass that thorugh the encoder\n",
    "         then use the encoding to pass through self.conv_net\n",
    "        '''\n",
    "        # obs = torch.tensor(obs)\n",
    "        if use_encoding:\n",
    "            obs = self.random_shift(process_obs(obs.copy()))\n",
    "            obs = self.CURL.encode(obs, detach=True)\n",
    "\n",
    "        # Dividing obs by 255 is handled in encoder forward (only needed for use_encoding=False)\n",
    "        if len(obs.shape) == 1:\n",
    "            obs = obs.unsqueeze(0)\n",
    "\n",
    "        # epsilon greedy for selecting which action to take\n",
    "        if random.random() > epsilon:\n",
    "            qs = self.conv_net(obs)\n",
    "            action = qs[0].argmax().data.item()\n",
    "        else:\n",
    "            action = random.randint(0, self.N_action-1)\n",
    "\n",
    "        return action\n",
    "\n",
    "def get_decay(epi_iter):\n",
    "    decay = math.pow(0.99, epi_iter)\n",
    "    return decay"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "1WWwX8wewOmy"
   },
   "outputs": [],
   "source": [
    "def main(aug=None, train_curve_filename_prefix=\"default_curl_cartpole\"):\n",
    "    env = gym.make('CartPole-v0')\n",
    "    env.reset()\n",
    "    max_epi_iter = 1000\n",
    "    max_MC_iter = 200\n",
    "    obs = env.render(mode='rgb_array')\n",
    "    obs_shape = obs.shape\n",
    "    shift_by = 20\n",
    "    random_shift = get_random_shift(*obs_shape[:2], shift_by)\n",
    "    cropped_obs_shape = (obs_shape[0] - shift_by, obs_shape[1] - shift_by, obs_shape[2])\n",
    "    encoder = PixelEncoder(obs_shape=cropped_obs_shape, input_channels=6, feature_dim=50, num_layers=2, num_filters=64, output_logits=False).to(device)\n",
    "    CURL_encoder = CURL(obs_shape=obs_shape, z_dim=50, batch_size=1, encoder=encoder, output_type=\"continuous\").to(device)\n",
    "    agent = Agent(act_dim=env.action_space.n, max_epi_num=max_epi_iter, max_epi_len=max_MC_iter, CURL=CURL_encoder, aug=aug, random_shift=random_shift)\n",
    "    train_curve = []\n",
    "    train_steps = []\n",
    "    \n",
    "    exploration_rate = 0.8\n",
    "    \n",
    "    for epi_iter in range(max_epi_iter):\n",
    "        random.seed()\n",
    "        env.reset()\n",
    "        obs = env.render(mode='rgb_array')\n",
    "        ''' For step 0, copy the same observation frames 2 times '''\n",
    "        # 400, 600, 6\n",
    "        stacked_obs = np.concatenate([obs, obs], -1)\n",
    "        returns = 0.0\n",
    "        steps = 0\n",
    "        for MC_iter in range(max_MC_iter):\n",
    "            exploration_rate *= get_decay(epi_iter)\n",
    "            exploration_rate = max(exploration_rate, 0.1)\n",
    "            action = agent.get_action(stacked_obs, exploration_rate)\n",
    "            next_obs, reward, done = take_action(env, action)\n",
    "            returns += reward * agent.gamma ** (MC_iter)\n",
    "           \n",
    "            if done or MC_iter >= max_MC_iter-1:\n",
    "                # agent.buffer.create_new_epi()\n",
    "                if MC_iter < max_MC_iter -1 :\n",
    "                    reward = 0 # - max_epi_iter\n",
    "                steps = MC_iter\n",
    "                \n",
    "                \n",
    "            ''' Stack 2 observation frames for input and 2 for output'''\n",
    "            stacked_next = np.concatenate([stacked_obs[:,:,-3:], next_obs], -1)\n",
    "            agent.remember(stacked_obs, action, reward, stacked_next)\n",
    "            obs = next_obs.copy()\n",
    "            stacked_obs = stacked_next.copy()\n",
    "            \n",
    "            if done:\n",
    "                break\n",
    "            \n",
    "        print('Episode', epi_iter, 'returns', returns, 'after', steps, 'timesteps')\n",
    "        if epi_iter % 1 == 0:\n",
    "            train_curve.append(returns)\n",
    "            train_steps.append(steps)\n",
    "        if epi_iter % 100 == 0:\n",
    "            print(f\"Saving at episode {epi_iter}\")\n",
    "            np.save(f'{train_curve_filename_prefix}_{max_MC_iter}MC_{max_epi_iter}_stacked_no_penalty_returns', np.array(train_curve))\n",
    "            np.save(f'{train_curve_filename_prefix}_{max_MC_iter}MC_{max_epi_iter}_stacked_no_penalty_steps', np.array(train_steps))\n",
    "        if agent.buffer.is_available():\n",
    "            for _ in range(1):\n",
    "                agent.train()\n",
    "    env.close()\n",
    "    np.save(f'{train_curve_filename_prefix}_{max_MC_iter}MC_{max_epi_iter}_stacked_returns_no_penalty', np.array(train_curve))\n",
    "    np.save(f'{train_curve_filename_prefix}_{max_MC_iter}MC_{max_epi_iter}_stacked_steps_no_penalty', np.array(train_steps))\n",
    "    print(train_curve)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "PpyPzumawOmz",
    "outputId": "e436d671-439f-408c-d4f2-ecac8fd5e2e5",
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "main()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "main(color_jitter, \"color_jitter\")\n",
    "main(random_elastic_transform, \"random_elastic_transform\")\n",
    "main(random_fisheye, \"random_fisheye\")\n",
    "main(random_color_equalize, \"random_color_equalize\")\n",
    "main(random_gaussian_blur, \"random_gaussian_blur\")\n",
    "main(random_gaussian_noise, \"random_gaussian_noise\")\n",
    "main(random_horizontal_flip, \"random_horizontal_flip\")\n",
    "main(random_color_invert, \"random_color_invert\")\n",
    "main(random_perspective_shift, \"random_perspective_shift\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "UkCfk2GjBNRl"
   },
   "source": [
    "Code references for DQN:\n",
    "\n",
    "https://github.com/taochenshh/dqn-pytorch\n",
    "\n",
    "https://github.com/transedward/pytorch-dqn (for sampling from replay buffer)\n",
    "\n",
    "CURL code: https://github.com/MishaLaskin/curl"
   ]
  }
 ],
 "metadata": {
  "colab": {
   "collapsed_sections": [],
   "name": "sample_gym_env_v3.ipynb",
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}