VAE.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.utils.data\n",
    "from torch import nn, optim, distributions\n",
    "from torch.nn import functional as F\n",
    "from torchvision import datasets, transforms\n",
    "from torchvision.utils import save_image\n",
    "import matplotlib.pyplot as plt\n",
    "from torch.utils.data import Dataset, TensorDataset, DataLoader\n",
    "%matplotlib inline\n",
    "\n",
    "from collections import deque\n",
    "import numpy as np\n",
    "from scipy.stats import mode"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n",
    "\n",
    "MNIST_train = torch.utils.data.DataLoader(\n",
    "    datasets.MNIST('../data', train=True, download=True,\n",
    "                   transform=transforms.ToTensor()),\n",
    "                   batch_size=128, shuffle=True)\n",
    "\n",
    "MNIST_test = torch.utils.data.DataLoader(\n",
    "    datasets.MNIST('../data', train=False, download=True,\n",
    "                   transform=transforms.ToTensor()),\n",
    "                   batch_size=128, shuffle=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_image(tensor, label=None):\n",
    "    '''Takes a [n, n] tensor, transform it to numpy and plot gray scale'''\n",
    "    tensor = tensor.cpu().numpy()\n",
    "    plt.imshow(tensor, cmap='gray', vmin=0, vmax=1)\n",
    "    if label is not None:\n",
    "        print('The correct label is {}'.format(label))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Success Unit\n",
    "\n",
    "The following block works well for MNIST. It use binary cross entropy loss. It also records the loss that is used in the following unit, which is $\\log p_\\theta(x|z)$ if $x=sigmoid(y)$ with $y\\sim\\mathcal{N}(\\mu_\\theta(z),\\sigma_\\theta(z))$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "class encoder(nn.Module):\n",
    "    '''this encoder is the recognition model'''\n",
    "    def __init__(self, input_dim, code_dim, hidden=None):\n",
    "        super(encoder, self).__init__()\n",
    "        if hidden is None:\n",
    "            hidden = [256, 256]\n",
    "        \n",
    "        #For new mean\n",
    "        self.fc1_1 = nn.Linear(input_dim, hidden[0])\n",
    "        self.fc1_2 = nn.Linear(hidden[0], hidden[1])\n",
    "        self.fc1_3 = nn.Linear(hidden[1], code_dim)\n",
    "        \n",
    "        #For new std\n",
    "        self.fc2_1 = nn.Linear(input_dim, hidden[0])\n",
    "        self.fc2_2 = nn.Linear(hidden[0], hidden[1])\n",
    "        self.fc2_3 = nn.Linear(hidden[1], code_dim)\n",
    "    \n",
    "    def forward(self, x):\n",
    "        '''\n",
    "        input x is an observation\n",
    "        returns mu and std such that p(z|x)~N(mu(x), std(x)). z is the hidden state\n",
    "        '''\n",
    "        \n",
    "        mu = F.relu(self.fc1_1(x))\n",
    "        mu = F.relu(self.fc1_2(mu))\n",
    "        mu = self.fc1_3(mu)\n",
    "        \n",
    "        log_std = F.relu(self.fc2_1(x))\n",
    "        log_std = F.relu(self.fc2_2(log_std))\n",
    "        log_std = self.fc2_3(log_std)\n",
    "        return mu, log_std\n",
    "    \n",
    "class decoder(nn.Module):\n",
    "    \n",
    "    def __init__(self, output_dim, code_dim, hidden=None):\n",
    "        super(decoder, self).__init__()\n",
    "        if hidden is None:\n",
    "            hidden = [256, 256]\n",
    "        \n",
    "        #for mean\n",
    "        self.fc1_1 = nn.Linear(code_dim, hidden[0])\n",
    "        self.fc1_2 = nn.Linear(hidden[0], hidden[1])\n",
    "        self.fc1_3 = nn.Linear(hidden[1], output_dim)\n",
    "        \n",
    "        #for std\n",
    "        self.fc2_1 = nn.Linear(code_dim, hidden[0])\n",
    "        self.fc2_2 = nn.Linear(hidden[0], hidden[1])\n",
    "        self.fc2_3 = nn.Linear(hidden[1], output_dim)\n",
    "    \n",
    "    def forward(self, z):\n",
    "        '''\n",
    "        input z is a hidden state\n",
    "        outputs mu and std\n",
    "        q(x|z)~sigmoid(N(mu(z), std(z))) because we need 0<=x<=1\n",
    "        '''\n",
    "        mu = F.relu(self.fc1_1(z))\n",
    "        mu = F.relu(self.fc1_2(mu))\n",
    "        mu = self.fc1_3(mu)\n",
    "        \n",
    "        log_std = F.relu(self.fc2_1(z))\n",
    "        log_std = F.relu(self.fc2_2(log_std))\n",
    "        log_std = self.fc2_3(log_std)\n",
    "        \n",
    "        return mu, log_std\n",
    "    \n",
    "class VAE1(nn.Module):\n",
    "    \n",
    "    def __init__(self, input_dim, code_dim, learning_rate):\n",
    "        super(VAE1, self).__init__()\n",
    "        self.encoder = encoder(input_dim, code_dim).to(device)\n",
    "        self.decoder = decoder(input_dim, code_dim).to(device)\n",
    "        self.encoder_optimizer = optim.Adam(self.encoder.parameters(), lr=learning_rate)\n",
    "        self.decoder_optimizer = optim.Adam(self.decoder.parameters(), lr=learning_rate)\n",
    "        \n",
    "    def forward(self, x, noise=False):\n",
    "        mu_z, log_std_z = self.encoder(x)\n",
    "        std_z = log_std_z.exp()\n",
    "        if noise:\n",
    "            z = distributions.Normal(mu_z, std_z).rsample()\n",
    "        else:\n",
    "            z = mu_z\n",
    "        mu_x, log_std_x = self.decoder(z)\n",
    "        std_x = log_std_x.exp()\n",
    "        if noise:\n",
    "            x = F.sigmoid(distributions.Normal(mu_x, std_x).rsample())\n",
    "        else:\n",
    "            x = F.sigmoid(mu_x)\n",
    "        return mu_z, log_std_z, mu_x, log_std_x, z, x\n",
    "    \n",
    "    def train(self, num_epoch, trainloader):\n",
    "        '''This function trains the VAE. Inputs are number of training epochs and the training loader'''\n",
    "        all_loss_log = []\n",
    "        loss_log = deque(maxlen = 50)\n",
    "        for i in range(1, num_epoch+1):\n",
    "            another_loss_log = deque(maxlen=50)\n",
    "            for j, (images, _) in enumerate(trainloader):\n",
    "                images = images.view(images.size(0), -1)\n",
    "                mu_z, log_std_z, mu_x, log_std_x, z, x = self.forward(images, noise=True)\n",
    "                #We will maximize the loss in paper, so we use the negative things here\n",
    "                KLD = -0.5*torch.sum(1 + 2*log_std_z - mu_z.pow(2) - (2*log_std_z).exp())\n",
    "                recon_loss = F.binary_cross_entropy(x, images, size_average=False)\n",
    "                Loss = KLD + recon_loss\n",
    "                self.encoder_optimizer.zero_grad()\n",
    "                self.decoder_optimizer.zero_grad()\n",
    "                Loss.backward()\n",
    "                self.encoder_optimizer.step()\n",
    "                self.decoder_optimizer.step()\n",
    "                loss_log.append(Loss.detach().cpu().numpy())\n",
    "                all_loss_log.append(loss_log[-1])\n",
    "                print('\\rEpisode {} batch group {}. Current average loss : {:.4f}'.format(i, j+1, np.mean(loss_log)/128), end='')\n",
    "                if (j+1) % 100 == 0:\n",
    "                    print('')\n",
    "            print('')\n",
    "        return all_loss_log"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 1 batch group 100. Current average loss : 206.8342\n",
      "Episode 1 batch group 200. Current average loss : 195.5303\n",
      "Episode 1 batch group 300. Current average loss : 185.2719\n",
      "Episode 1 batch group 400. Current average loss : 176.8340\n",
      "Episode 1 batch group 469. Current average loss : 168.2712\n",
      "Episode 2 batch group 100. Current average loss : 156.7263\n",
      "Episode 2 batch group 200. Current average loss : 151.1019\n",
      "Episode 2 batch group 300. Current average loss : 146.1395\n",
      "Episode 2 batch group 400. Current average loss : 141.6671\n",
      "Episode 2 batch group 469. Current average loss : 138.1352\n",
      "Episode 3 batch group 100. Current average loss : 135.0700\n",
      "Episode 3 batch group 200. Current average loss : 133.1129\n",
      "Episode 3 batch group 300. Current average loss : 129.8708\n",
      "Episode 3 batch group 400. Current average loss : 127.6247\n",
      "Episode 3 batch group 469. Current average loss : 125.0645\n",
      "Episode 4 batch group 100. Current average loss : 124.2102\n",
      "Episode 4 batch group 200. Current average loss : 122.6800\n",
      "Episode 4 batch group 300. Current average loss : 121.9622\n",
      "Episode 4 batch group 400. Current average loss : 121.1182\n",
      "Episode 4 batch group 469. Current average loss : 118.9896\n",
      "Episode 5 batch group 100. Current average loss : 119.7470\n",
      "Episode 5 batch group 200. Current average loss : 118.0941\n",
      "Episode 5 batch group 300. Current average loss : 117.8290\n",
      "Episode 5 batch group 400. Current average loss : 117.6524\n",
      "Episode 5 batch group 469. Current average loss : 115.9536\n",
      "Episode 6 batch group 100. Current average loss : 116.2331\n",
      "Episode 6 batch group 200. Current average loss : 116.0139\n",
      "Episode 6 batch group 300. Current average loss : 113.9386\n",
      "Episode 6 batch group 400. Current average loss : 113.9896\n",
      "Episode 6 batch group 469. Current average loss : 114.2086\n",
      "Episode 7 batch group 100. Current average loss : 113.6303\n",
      "Episode 7 batch group 200. Current average loss : 113.0097\n",
      "Episode 7 batch group 300. Current average loss : 112.9860\n",
      "Episode 7 batch group 400. Current average loss : 111.9284\n",
      "Episode 7 batch group 469. Current average loss : 111.6505\n",
      "Episode 8 batch group 100. Current average loss : 111.5587\n",
      "Episode 8 batch group 200. Current average loss : 111.3129\n",
      "Episode 8 batch group 300. Current average loss : 111.0965\n",
      "Episode 8 batch group 400. Current average loss : 110.5960\n",
      "Episode 8 batch group 469. Current average loss : 109.6024\n",
      "Episode 9 batch group 100. Current average loss : 110.3734\n",
      "Episode 9 batch group 200. Current average loss : 110.3057\n",
      "Episode 9 batch group 300. Current average loss : 109.5564\n",
      "Episode 9 batch group 400. Current average loss : 109.7584\n",
      "Episode 9 batch group 469. Current average loss : 107.9646\n",
      "Episode 10 batch group 100. Current average loss : 108.7177\n",
      "Episode 10 batch group 200. Current average loss : 107.7477\n",
      "Episode 10 batch group 300. Current average loss : 108.3346\n",
      "Episode 10 batch group 400. Current average loss : 108.1180\n",
      "Episode 10 batch group 469. Current average loss : 108.0164\n"
     ]
    }
   ],
   "source": [
    "MNIST_VAE1 = VAE1(input_dim=28*28, code_dim=20, learning_rate=5e-4)\n",
    "MNIST_losses = MNIST_VAE1.train(num_epoch = 10, trainloader=MNIST_train)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The correct label is 1\n"
     ]
    },
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAA2oAAAEZCAYAAADmAtZNAAAABHNCSVQICAgIfAhkiAAAAAlwSFlzAAALEgAACxIB0t1+/AAAADl0RVh0U29mdHdhcmUAbWF0cGxvdGxpYiB2ZXJzaW9uIDMuMC4yLCBodHRwOi8vbWF0cGxvdGxpYi5vcmcvOIA7rQAAFtZJREFUeJzt3V2sXWd5J/D/4w8wJh/EARKTOMQJZDRQFDNYaCSqEaiiML0xXDRqLoaMVMZcFKlIczEIIRUJRkKjwkyFEFIKUVLRUipRhqiqJiWoEhlpRHAQCgmBEFCmdWLFQflwnBA79nnnwjsa17HPWufsvc9+zzm/nxSdc/Z6vNajFe/H+3/W2vut1loAAADox5ZFNwAAAMC/JKgBAAB0RlADAADojKAGAADQGUENAACgM4IaAABAZwQ1AACAzghqAAAAnRHUAAAAOiOoAQAAdGbbNH+4qj6Y5M+SbE3y1dba5wfq2zTHA/rUWqtF93Culcwnswk2rF+31t6w6CbO5rUTkIx77VStre75X1Vbkzyc5P1JDif5YZKbW2s/XebPGDawAfUW1FY6n8wm2LDua63tX3QTL/PaCXjZmNdO09z6+O4kj7TWftVaO5nkr5McmGJ/ALNiPgE9MpuA0aYJalcl+eezfj48eQxg0cwnoEdmEzDaNO9RO9/luldcnq+qg0kOTnEcgJUanE9mE7AAXjsBo00T1A4n2XPWz1cnefzcotbarUluTdxnDayZwflkNgEL4LUTMNo0tz7+MMlbq2pvVb0qyR8kuXM2bQFMxXwCemQ2AaOt+opaa+1UVX08yV058xGzt7XWHpxZZwCrZD4BPTKbgJVY9cfzr+pgLt/DhtTbx/OvlNkEG1ZXH8+/GuYTbEzz/nh+AAAA5kBQAwAA6IygBgAA0BlBDQAAoDOCGgAAQGcENQAAgM4IagAAAJ0R1AAAADojqAEAAHRGUAMAAOiMoAYAANAZQQ0AAKAzghoAAEBnBDUAAIDOCGoAAACdEdQAAAA6s23RDQDAWtmyZTa/n2ytzaQGAC7EFTUAAIDOCGoAAACdEdQAAAA6I6gBAAB0RlADAADojKAGAADQGUENAACgM4IaAABAZyx4zVx8+tOfHqz57Gc/O1hz4MCBwZo777xzVE/A+jZmseobbrhh2e2f+9znBvdx1VVXDdZ89atfHaz5+te/Plhz4sSJwRqgf2Pm0969e5fd/pGPfGQmx/nSl740WHP06NHBGhbPFTUAAIDOCGoAAACdEdQAAAA6I6gBAAB0RlADAADojKAGAADQGUENAACgM4IaAABAZyx4zaq8+c1vXnb7LbfcMriPpaWlWbUDbAJbt24drHnf+9637PZ3vOMdg/vYuXPnYM2Y/VTVYA2wMezYsWOw5qabblp2+wc+8IHBfbz00kuDNd/85jcHa5588snBmtbaYA3zNVVQq6pHkzyX5HSSU621/bNoCmBa5hPQI7MJGGsWV9Te11r79Qz2AzBr5hPQI7MJGOQ9agAAAJ2ZNqi1JP9QVfdV1cFZNAQwI+YT0COzCRhl2lsf39Nae7yq3pjku1X1s9ba988umAwhgwhYa8vOJ7MJWBCvnYBRprqi1lp7fPL1aJJvJ3n3eWpuba3t92ZZYC0NzSezCVgEr52AsVYd1KrqtVV18cvfJ/ndJA/MqjGA1TKfgB6ZTcBKTHPr4xVJvj1ZJ2Zbkr9qrf2vmXQFMB3zCeiR2QSMtuqg1lr7VZIbZ9gL68h111031XaYJ/NpYxqz+OoVV1yx7PZLL710cB8nTpwYrPnZz342WHPy5MnBGjYXs2njWlpaGqy58cbl/9fv3bt3cB8PP/zwYM3x48cHayxmvT74eH4AAIDOCGoAAACdEdQAAAA6I6gBAAB0RlADAADojKAGAADQGUENAACgM4IaAABAZ1a94DUArKWqGqx5wxvesOz2iy++eCbHuffeewdrxiyAC2wMr3nNawZrrrnmmmW3X3LJJYP72Llz52DNk08+OVjD+uCKGgAAQGcENQAAgM4IagAAAJ0R1AAAADojqAEAAHRGUAMAAOiMoAYAANAZQQ0AAKAzFrxmVd71rnctugWAV7juuuuW3b5t2/A/e621wRoLygIrtWfPnmW3j5lPW7duHaw5efLk6J7omytqAAAAnRHUAAAAOiOoAQAAdEZQAwAA6IygBgAA0BlBDQAAoDOCGgAAQGeso8aqHDhwYOp93H///YM1d99999THATaGHTt2DNZcf/31y26vqsF9jFmD6NixY4M1wOaxa9euwZpLLrlk2e1j5tPp06cHa5aWlgZrWB9cUQMAAOiMoAYAANAZQQ0AAKAzghoAAEBnBDUAAIDOCGoAAACdEdQAAAA6I6gBAAB0xoLXLMyYRWVfeOGFNegEWA+GFrNOkssvv3zZ7WMWlD116tRgzUsvvTRYA2web3vb2wZrtm/fvuz2MfPpxIkTo3ti/Ru8olZVt1XV0ap64KzHdlXVd6vqF5Ovl823TYBXMp+AHplNwCyMufXx9iQfPOexTyb5XmvtrUm+N/kZYK3dHvMJ6M/tMZuAKQ0Gtdba95M8dc7DB5LcMfn+jiQfmnFfAIPMJ6BHZhMwC6v9MJErWmtHkmTy9Y2zawlgKuYT0COzCViRuX+YSFUdTHJw3scBWAmzCeiV+QQkq7+i9kRV7U6SydejFypsrd3aWtvfWtu/ymMBrMSo+WQ2AWvMaydgRVYb1O5Mcsvk+1uSfGc27QBMzXwCemQ2ASsy5uP5v5Hk/yT5V1V1uKr+MMnnk7y/qn6R5P2TnwHWlPkE9MhsAmZh8D1qrbWbL7Dpd2bcC5346Ec/Olizb9++qY/z5S9/eep9sLmZTxvHmIVe3/72tw/WDC0oO8bTTz89WGPBa5ZjNm0sY+bTjTfeOFizZctqb2T7/5599tnBmqWlpamPQx+m/xsDAADATAlqAAAAnRHUAAAAOiOoAQAAdEZQAwAA6IygBgAA0BlBDQAAoDOCGgAAQGcGF7xm8zl48OBgzc6dO6c+zhNPPDH1PoCNYcyCssePHx+saa1NfZwXX3xxsGbMfoCNYczz/eGHH16T43jttLm4ogYAANAZQQ0AAKAzghoAAEBnBDUAAIDOCGoAAACdEdQAAAA6I6gBAAB0RlADAADojAWvAVi4MQu97tq1a7Bm+/btUx/nxIkTgzUWvAbOtnfv3sGaLVumvz5y8cUXT70P1g9X1AAAADojqAEAAHRGUAMAAOiMoAYAANAZQQ0AAKAzghoAAEBnBDUAAIDOWEcNgHXhhhtuGKyZxTpqx44dG90TwFizWEfthRdemEEnrBeuqAEAAHRGUAMAAOiMoAYAANAZQQ0AAKAzghoAAEBnBDUAAIDOCGoAAACdEdQAAAA6Y8FrABZu69atgzX79u0brBlaUHZpaWlwH0eOHBmsOX369GANsDFU1WDNlVdeOfV+WmuD+zh8+PBgDRvH4BW1qrqtqo5W1QNnPfaZqnqsqn48+e/35tsmwCuZT0CPzCZgFsbc+nh7kg+e5/H/3lrbN/nv72fbFsAot8d8Avpze8wmYEqDQa219v0kT61BLwArYj4BPTKbgFmY5sNEPl5V908u7182s44Apmc+AT0ym4DRVhvUvpLk+iT7khxJ8oULFVbVwao6VFWHVnksgJUYNZ/MJmCNee0ErMiqglpr7YnW2unW2lKSP0/y7mVqb22t7W+t7V9tkwBjjZ1PZhOwlrx2AlZqVUGtqnaf9eOHkzxwoVqAtWQ+AT0ym4CVGlxHraq+keS9SV5fVYeT/EmS91bVviQtyaNJPjbHHgHOy3wCemQ2AbMwGNRaazef5+GvzaEXgBUxnzaObdsG/znK1VdfPfVxxixUfffddw/WjFk4m83LbNpYxixEff311w/WzGLB60OHht+2OGY/rA/TfOojAAAAcyCoAQAAdEZQAwAA6IygBgAA0BlBDQAAoDOCGgAAQGcENQAAgM4IagAAAJ0ZXmEUADowZpHpoYVexyx4/dRTT43uCdj4hhaqTpKdO3euQSfJc889tybHoQ+uqAEAAHRGUAMAAOiMoAYAANAZQQ0AAKAzghoAAEBnBDUAAIDOCGoAAACdEdQAAAA6Y8FrABZu+/btgzVjFqseWvD61KlTg/vYtWvXYM2WLcO/5xyzQDfQv61btw7WXHLJJVMfZ2h+Jcnu3bsHa8Ys0D3mWCyeK2oAAACdEdQAAAA6I6gBAAB0RlADAADojKAGAADQGUENAACgM4IaAABAZwQ1AACAzljwepPZt2/fYM2111479XEeeeSRwZoHHnhg6uMAG8OOHTsGa2axoOzJkycHa5555pmpjwNsHNu3bx+smcV8On369GDNmMWst2wZvg4z5lgsnitqAAAAnRHUAAAAOiOoAQAAdEZQAwAA6IygBgAA0BlBDQAAoDOCGgAAQGeso7bJjFkf6Pnnnx+sufzyy5fdftFFFw3u49JLLx2seeyxxwZrgPVv27bhf45e/epXD9a01pbdfurUqcF9jFmnaOg4wMYxZh21MfNpaWlpqu1Jct111w3WjJlhrA+DV9Sqak9V/WNVPVRVD1bVH08e31VV362qX0y+Xjb/dgHOMJuAXplPwCyMufXxVJL/3Fr710n+bZI/qqq3Jflkku+11t6a5HuTnwHWitkE9Mp8AqY2GNRaa0daaz+afP9ckoeSXJXkQJI7JmV3JPnQvJoEOJfZBPTKfAJmYUUfJlJV1yZ5Z5IfJLmitXYkOTOQkrxx1s0BjGE2Ab0yn4DVGv1hIlV1UZJvJflEa+3Y2DcqVtXBJAdX1x7A8swmoFfmEzCNUVfUqmp7zgyav2yt/e3k4Seqavdk++4kR8/3Z1trt7bW9rfW9s+iYYCXmU1Ar8wnYFpjPvWxknwtyUOttS+etenOJLdMvr8lyXdm3x7A+ZlNQK/MJ2AWxtz6+J4k/yHJT6rqx5PHPpXk80n+pqr+MMk/Jfn9+bQIcF5mE9Ar8wmY2mBQa6397yQXuqn6d2bbDvP26KOPDtY8+eSTgzXXXHPNstuvvPLKwX3s2bNnsOanP/3pYA2bk9m0sYxZUPb48eODNa973euW3X769OnBfYxZfNuCsizHfNpYxjzfn3766cGaK664YtntY+bTsWPHBmvGzLBTp04N1rB4K/rURwAAAOZPUAMAAOiMoAYAANAZQQ0AAKAzghoAAEBnBDUAAIDOCGoAAACdEdQAAAA6M7wiHqzCmIW1f/7zn8+/EWBdeOmllwZrHn/88cGaq6++etntYxaC/c1vfjNY01obrAE2hpMnTw7WPPjgg4M1b3nLW5bdvrS0NLiPZ555ZrDGYtYbhytqAAAAnRHUAAAAOiOoAQAAdEZQAwAA6IygBgAA0BlBDQAAoDOCGgAAQGcENQAAgM5Y8Jq5eOyxxwZrxiyKDWwOx44dG6w5dOjQYM2b3vSmZbffe++9g/v45S9/OVgzZmFaYGMYs+D1XXfdNVgztOD14cOHB/dxzz33DNaYTxuHK2oAAACdEdQAAAA6I6gBAAB0RlADAADojKAGAADQGUENAACgM4IaAABAZwQ1AACAzlRrbe0OVrV2BwPWTGutFt3DNMym9WHLluHfLb7qVa9advvWrVsH9/HCCy8M1qzlv51M5b7W2v5FNzEN82l9qBr+Z3Dnzp3Lbt+xY8fgPp599tnBmlOnTg3WsHhjXju5ogYAANAZQQ0AAKAzghoAAEBnBDUAAIDOCGoAAACdEdQAAAA6I6gBAAB0RlADAADozLahgqrak+QvklyZZCnJra21P6uqzyT5T0menJR+qrX29/NqFOBsZtPms7S0NFjz4osvrkEnsDzzafNpbXhd8ueff36q7Ww+NfQXq6p2J9ndWvtRVV2c5L4kH0pyU5LjrbU/HX2wquG/xcC601qrtT6m2QSMcF9rbf9aH9R8AoaMee00eEWttXYkyZHJ989V1UNJrpq+PYDVM5uAXplPwCys6D1qVXVtkncm+cHkoY9X1f1VdVtVXTbj3gBGMZuAXplPwGqNDmpVdVGSbyX5RGvtWJKvJLk+yb6c+a3RFy7w5w5W1aGqOjSDfgH+BbMJ6JX5BExj8D1qSVJV25P8XZK7WmtfPM/2a5P8XWvttwb24z5r2IAW8R61xGwCBi3kPWqJ+QQsb8xrp8EralVVSb6W5KGzB83kjbIv+3CSB1bTJMBqmE1Ar8wnYBbGfOrjbye5J8lPcuYjZpPkU0luzplL9y3Jo0k+Nnnz7HL78lsh2IAW9KmPZhMwZFGf+mg+Acsa89pp1K2Ps2LYwMa0qFsfZ8Vsgg1rYbc+zor5BBvTTG59BAAAYG0JagAAAJ0R1AAAADojqAEAAHRGUAMAAOiMoAYAANAZQQ0AAKAzghoAAEBnBDUAAIDOCGoAAACdEdQAAAA6I6gBAAB0RlADAADojKAGAADQGUENAACgM4IaAABAZwQ1AACAzmxb4+P9Osn/Pevn108eWy/0O1/6na959fvmOexzrZ07mxL/f+dNv/Ol3zM24nzy/3a+9Dtf+j1j1Gyq1tocjj1OVR1qre1fWAMrpN/50u98rbd+F229nS/9zpd+52u99btI6+1c6Xe+9Dtfi+7XrY8AAACdEdQAAAA6s+igduuCj79S+p0v/c7Xeut30dbb+dLvfOl3vtZbv4u03s6VfudLv/O10H4X+h41AAAAXmnRV9QAAAA4x8KCWlV9sKp+XlWPVNUnF9XHWFX1aFX9pKp+XFWHFt3Puarqtqo6WlUPnPXYrqr6blX9YvL1skX2eLYL9PuZqnpsco5/XFW/t8gez1ZVe6rqH6vqoap6sKr+ePJ4l+d4mX67Pce9MJtmy2yaL7NpczGfZst8mh+zaUZ9LeLWx6ramuThJO9PcjjJD5Pc3Fr76Zo3M1JVPZpkf2uty7UfqurfJTme5C9aa781eey/JXmqtfb5yUC/rLX2XxbZ58su0O9nkhxvrf3pIns7n6ranWR3a+1HVXVxkvuSfCjJf0yH53iZfm9Kp+e4B2bT7JlN82U2bR7m0+yZT/NjNs3Goq6ovTvJI621X7XWTib56yQHFtTLhtBa+36Sp855+ECSOybf35Ezf+G6cIF+u9VaO9Ja+9Hk++eSPJTkqnR6jpfpl+WZTTNmNs2X2bSpmE8zZj7Nj9k0G4sKalcl+eezfj6cDk7GgJbkH6rqvqo6uOhmRrqitXYkOfMXMMkbF9zPGB+vqvsnl/e7uBx+rqq6Nsk7k/wg6+Acn9Nvsg7O8QKZTWuj++fNeXT/vDGbNjzzaW10/9w5j66fO2bT6i0qqNV5Huv94yff01r7N0n+fZI/mlx+Zra+kuT6JPuSHEnyhcW280pVdVGSbyX5RGvt2KL7GXKefrs/xwtmNnE+3T9vzKZNwXzifLp+7phN01lUUDucZM9ZP1+d5PEF9TJKa+3xydejSb6dM7cg9O6JyT23L997e3TB/SyrtfZEa+10a20pyZ+ns3NcVdtz5sn7l621v5083O05Pl+/vZ/jDphNa6Pb58359P68MZs2DfNpbXT73Dmfnp87ZtP0FhXUfpjkrVW1t6peleQPkty5oF4GVdVrJ28sTFW9NsnvJnlg+T/VhTuT3DL5/pYk31lgL4NefuJOfDgdneOqqiRfS/JQa+2LZ23q8hxfqN+ez3EnzKa10eXz5kJ6ft6YTZuK+bQ2unzuXEivzx2zaUZ9LWrB6zrz8Zb/I8nWJLe11v7rQhoZoaquy5nfBCXJtiR/1Vu/VfWNJO9N8vokTyT5kyT/M8nfJLkmyT8l+f3WWhdvQr1Av+/NmUvLLcmjST728n3Mi1ZVv53kniQ/SbI0efhTOXP/cnfneJl+b06n57gXZtNsmU3zZTZtLubTbJlP82M2zaivRQU1AAAAzm9hC14DAABwfoIaAABAZwQ1AACAzghqAAAAnRHUAAAAOiOoAQAAdEZQAwAA6IygBgAA0Jn/BwSwSBE4/X4xAAAAAElFTkSuQmCC\n",
      "text/plain": [
       "<Figure size 1080x540 with 3 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "image, label = iter(MNIST_train).next()\n",
    "with torch.no_grad():\n",
    "    mu_z1, log_std_z1, mu_x1, log_std_x1, z1, x1 = MNIST_VAE1(image[0,:,:,:].view(-1), noise=False)\n",
    "    mu_z2, log_std_z2, mu_x2, log_std_x2, z2, x2 = MNIST_VAE1(image[0,:,:,:].view(-1), noise=True)\n",
    "x1 = x1.view(28,28)\n",
    "x2 = x2.view(28,28)\n",
    "plt.figure(figsize=(15,7.5))\n",
    "plt.subplot(1,3,1)\n",
    "plot_image(tensor=image[0,0,:,:], label=label[0])\n",
    "plt.subplot(1,3,2)\n",
    "plot_image(x1)\n",
    "plt.subplot(1,3,3)\n",
    "plot_image(x2)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Test on gym, pong"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "import gym\n",
    "env = gym.make('PongDeterministic-v4')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "def process_image(state):\n",
    "    graph = state.copy()\n",
    "    graph = np.mean(graph, axis=2, keepdims=False)\n",
    "    graph = graph[34:194,:]\n",
    "    graph = np.where(graph==mode(graph, axis=None)[0].item(), 0, 1)\n",
    "    return graph"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "n_episode = 30\n",
    "max_t = 1000\n",
    "t_max = 500\n",
    "scores = []\n",
    "recorder = []\n",
    "for i in range(1, n_episode+1):\n",
    "    state = env.reset()\n",
    "    score = 0\n",
    "    for t in range(max_t):\n",
    "        recorder.append(process_image(state))\n",
    "        action = np.random.randint(0,6)\n",
    "        env.render()\n",
    "        next_state, reward, done, _ = env.step(action)\n",
    "        state = next_state\n",
    "        score += reward\n",
    "        if done:\n",
    "            break \n",
    "    scores.append(score)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "26865"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(recorder)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset = torch.stack([torch.tensor(image, dtype=torch.float32) for image in recorder])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([26865, 160, 160])"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dataset.size()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "Pong_train = DataLoader(TensorDataset(dataset, torch.zeros((dataset.size(0), 1))), batch_size=128, shuffle=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 1 batch group 100. Current average loss : 421.61145\n",
      "Episode 1 batch group 200. Current average loss : 401.7343\n",
      "Episode 1 batch group 210. Current average loss : 397.3952\n",
      "Episode 2 batch group 100. Current average loss : 370.3159\n",
      "Episode 2 batch group 200. Current average loss : 332.2101\n",
      "Episode 2 batch group 210. Current average loss : 324.0155\n",
      "Episode 3 batch group 100. Current average loss : 230.2783\n",
      "Episode 3 batch group 200. Current average loss : 165.2853\n",
      "Episode 3 batch group 210. Current average loss : 159.7938\n",
      "Episode 4 batch group 100. Current average loss : 125.7207\n",
      "Episode 4 batch group 200. Current average loss : 109.0150\n",
      "Episode 4 batch group 210. Current average loss : 107.0340\n",
      "Episode 5 batch group 100. Current average loss : 96.5336\n",
      "Episode 5 batch group 200. Current average loss : 90.0889\n",
      "Episode 5 batch group 210. Current average loss : 89.0262\n",
      "Episode 6 batch group 100. Current average loss : 83.6437\n",
      "Episode 6 batch group 200. Current average loss : 79.6485\n",
      "Episode 6 batch group 210. Current average loss : 79.3311\n",
      "Episode 7 batch group 100. Current average loss : 75.0403\n",
      "Episode 7 batch group 200. Current average loss : 72.7627\n",
      "Episode 7 batch group 210. Current average loss : 72.2616\n",
      "Episode 8 batch group 100. Current average loss : 69.0097\n",
      "Episode 8 batch group 200. Current average loss : 67.8479\n",
      "Episode 8 batch group 210. Current average loss : 67.1136\n",
      "Episode 9 batch group 100. Current average loss : 65.2223\n",
      "Episode 9 batch group 200. Current average loss : 63.9433\n",
      "Episode 9 batch group 210. Current average loss : 63.7060\n",
      "Episode 10 batch group 100. Current average loss : 60.5650\n",
      "Episode 10 batch group 200. Current average loss : 60.7239\n",
      "Episode 10 batch group 210. Current average loss : 60.1614\n"
     ]
    }
   ],
   "source": [
    "Pong_VAE1 = VAE1(input_dim=160*160, code_dim=20, learning_rate=1e-3)\n",
    "Pong_losses = Pong_VAE1.train(num_epoch = 10, trainloader=Pong_train)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAABIcAAAFrCAYAAABLz+nFAAAABHNCSVQICAgIfAhkiAAAAAlwSFlzAAALEgAACxIB0t1+/AAAADl0RVh0U29mdHdhcmUAbWF0cGxvdGxpYiB2ZXJzaW9uIDMuMC4yLCBodHRwOi8vbWF0cGxvdGxpYi5vcmcvOIA7rQAAIABJREFUeJzt3X+MpHd9H/D3Z2/P57PBsjGCgs8pJrJIE5Q0CPEjiSoESWMSZPNHojqijZtQWU3ThKRFAYrUKH8lUaKQRE2oXELsVAiSkh9YqPmBCIhGKiT8COBgHFwg9sGB+REbzmf7bm+//WPm4XbXu7N7O7M7z8zzekmr3ZmdnfnsMzvzvvs8n+f7VGstAAAAAAzTyrwLAAAAAGB+NIcAAAAABkxzCAAAAGDANIcAAAAABkxzCAAAAGDANIcAAAAABkxzCAAAAGDADqw5VFU3VNU9VXVvVb32oB4HgMUjIwCYRE4AHK5qrc3+TquOJPn7JN+X5GSSv0nyI621T8z8wQBYKDICgEnkBMDhO6jJoeclube19unW2tkkb0ty0wE9FgCLRUYAMImcADhkqwd0v9ckuX/D5ZNJnr/Tjatq9uNLAEuitVbzrmHGLiojEjkBMImckBMAk+wlJw6qObTdA296w66qW5PcekCPD0B/7ZoRiZwAGDA5AXDIDqo5dDLJtRsun0jy+Y03aK3dluS2RKcfYGB2zYhETgAMmJwAOGQHtebQ3yS5vqquq6pLktyc5M4DeiwAFouMAGASOQFwyA5kcqi1tlZV/zHJnyc5kuTNrbW/O4jHAmCxyAgAJpETAIfvQE5lf9FFGAMF2NESLjR60eQEwM7khJwAmGQvOXFQh5UBAAAAsAA0hwAAAAAGTHMIAAAAYMA0hwAAAAAGTHMIAAAAYMA0hwAAAAAGTHMIAAAAYMA0hwAAAAAGTHMIAAAAYMA0hwAAAAAGTHMIAAAAYMA0hwAAAAAGTHMIAAAAYMA0hwAAAAAGTHMIAAAOSVWlquZdBgA9tbKykpWVw2/VaA4BAAAADJjmEAAAAMCAaQ4BAAAADNjqvAsAAIChaK3NuwQAemxeOWFyCAAAAGDANIcAAAAABkxzCAAAAGDArDkEAAAA0APWHAIAAADg0JkcAgAAmNLKyuT97uvr64dUCcDFMzkEAAAAMGCaQwAA7KqqUlXzLgOAnlpZWdl1go7+8swBAAAADJg1hwAA2NW8zp4Ci8JrhKHqpkq9BhabySEAAACAAdMcAgBgV9YcAmA7rTVTQ0tAcwgAAABgwPbdHKqqa6vqPVV1d1X9XVW9anz9k6rqXVX1qfHnq2ZXLgCLQk7AcujOPmPPMLO2bDnRvUZ2+oBltbq6mtXVVX/rC672++RV1dOSPK219uGqemKSDyV5eZJ/m+SrrbVfqqrXJrmqtfaaXe7LXxDADlprC3kch5yA5dCdlnh9fX3OlbATOSEnYJ5WV0fnuVpbW5tzJexkLzmx78mh1tqp1tqHx19/PcndSa5JclOSO8Y3uyOjN3gABkZOwHJYX1/XGOJALFtOWJeLoVpbW9MYWgIzWXOoqp6R5DuTfCDJU1trp5LRG36Sp8ziMQBYXHICgEnkBMB8rU57B1X1hCR/mORnWmtf22u3vKpuTXLrtI8PQL/JCQAmWZacsNYKsMj2veZQklTV0STvTPLnrbVfG193T5IXtdZOjY8jfm9r7Vm73I93UoAdLOpaEomcADgMckJOAExyoGsO1ail/ztJ7u7eyMfuTHLL+Otbkrxjv48BwOKSEwBMIicA+mOas5V9T5L/k+TjSbpVCv9LRscJ/0GSb0pyX5Ifbq19dZf70ukH2MGi7hGWEwCHQ07ICYBJ9pITUx1WNivezAF2tqj/6J8lOQGwMzkhJwAmOdDDygAAAABYfJpDAAAAAAOmOQQAAAAwYJpDAAAAAAOmOQQAAAAwYJpDAAAAAAOmOQQAAAAwYJpDAAAAAAOmOQQAAAAwYJpDAAAAAAOmOQQAAAAwYJpDAAAAAAOmOQQAAAAwYJpDAAAAAAOmOQQAAAAwYJpDAAAAAAOmOQQAAAAwYJpDAAAAAAOmOQQAAAAwYJpDAAAAAAOmOQQAAAAwYJpDAAAAAAOmOQQAAAAwYJpDAAAAAAOmOQQAAAAwYJpDAAAAAAOmOQQAAAAwYJpDAAAAAAOmOQQAAAAwYJpDAAAAAAOmOQQAAAAwYJpDAAAAAAM2dXOoqo5U1Ueq6p3jy9dV1Qeq6lNV9ftVdcn0ZQKwqOQEAJPICYD5m8Xk0KuS3L3h8i8neUNr7fok/5jklTN4DAAWl5wAYBI5ATBnUzWHqupEkh9M8qbx5Ury4iRvH9/kjiQvn+YxAFhccgKASeQEQD9MOzn060l+Lsn6+PLVSR5sra2NL59Mcs2UjwHA4pITAEwiJwB6YN/Noap6WZIHWmsf2nj1NjdtO/z8rVX1war64H5rAKC/5AQAk8gJgP5YneJnvzvJjVX1A0kuTXJFRp3/K6tqddztP5Hk89v9cGvttiS3JUlVbfuGD8BCkxMATCInAHpi35NDrbXXtdZOtNaekeTmJH/ZWntFkvck+aHxzW5J8o6pqwRg4cgJACaREwD9MYuzlW31miT/qaruzeiY4d85gMcAYHHJCQAmkRMAh6xam/8EpjFQgJ211rZbf2FQ5ATAzuSEnACYZC85cRCTQwAAAAAsCM0hAAAAgAHTHAIAAAAYMM0hAAAAgAHTHAIAAAAYMM0hAAAAgAHTHAIA6LmqStXgz1YOABwQzSEAAACAAdMcAgDoOZNDAEwiJ5iW5hAAAADAgK3OuwAAgP04cuRIkqS1liRZX1+fZzkHwl5ggP3b+h7a5QXweCaHAAAAAAbM5BAAsNC6CaJlnBzq9nLb2w2wf90E0TK/ly7z78bhMDkEAAAAMGAmhwCAhdTtCV7GiSEAZsdUDezO5BAAAADAgJkcAgAW2srKaF/X+fPnZ3q/l19+eZILk0mPPPLITO8fgMV26aWXJrmQE2fPnp1nOTAVk0MAAAAAA2ZyCABYSGtrawd6/9YyAlhsB73WUJcT1jRiGZgcAgAAABgwk0MAAACwT93ZM2GRmRwCAAAAGDCTQwAA21hdHf0z6aDXNprGxr3V1rwAOFxHjhxJMvuzZcI8mBwCAAAAGLDqw16mqpp/EQA91Vob/IHscgK21+21Tuy5HjI5ISdgJysrF+ZBnIVzuPaSEyaHAAAAAAbMmkMAQO8cPXr0cdedO3duDpX028ZpoW79oT5MhQMctG5ycuPaa31eI25eNk4LyQkmMTkEAAAAMGAmhwCA3un2am7cIzxrW+970fekLnr9ABfjMN7zlm3SZll+Dw6GySEAAACAATM5BAD0zmFMDgHAJCZtGBKTQwAAAAADNlVzqKqurKq3V9Unq+ruqnphVT2pqt5VVZ8af75qVsUCsFiWKSeqKlWVlZWVrKysfONy98Fsddt5lrY+Z91jtNbsHYY5Waac4HB17+Xde/gs3sd3ygkYgmn/0n8jyZ+11r4lyXckuTvJa5O8u7V2fZJ3jy8DMExyAoBJ5ARAD9R+O6xVdUWSjyZ5ZttwJ1V1T5IXtdZOVdXTkry3tfasXe7L7jqAHbTWFnIsZdlyYuuew635OcvJkyNHjmy6fP78+U2Xu0ml7vP6+vrMHvuwdNuz225bt99201jTbuOdJrx2246Tzlaz2xnPdvs9D9LFnmVnp9t3f4/d73Lu3Lltf65jCmuygzj7kZzoR04cpq2ZtJ/3r76bR+0H8X62230uw3voXv8emZ+95MQ0k0PPTPKlJL9bVR+pqjdV1eVJntpaOzUu4FSSp0zxGAAsLjkBwCRyAqAnpjlb2WqS5yT5qdbaB6rqN3IRI59VdWuSW6d4fAD6baly4oUvfGGS5Ctf+UqS5PLLL0+SPPbYY0mSu+66a2aPtdtew2VY42inPaXd9ZdcckmSzVMqO22Xve5d3un7095vX1x66aVJLtR79uzZi/r53bbDbnuCD3N7Ldpzw46WKie6Kbud/i5nOU0xhJzobH29d5e77b22tjazxzqI95Q+vU9t/Rud9YTPXv8u+7RN+mhe22mayaGTSU621j4wvvz2jN7cvzge/8z48wPb/XBr7bbW2nNba8+dogYA+ktOADCJnADoiX1PDrXWvlBV91fVs1pr9yR5SZJPjD9uSfJL48/vmEmlACyUZcuJL3/5y0mSJz7xiUmSr33ta0mS1dVphnCns8h73rauo7RVH363vdSw2969rXu6Z/17nThx4htfd5NC3d9mZ1brHu13Mgt2smw5cdCv9/3oQw37tddpz41TUn38fbv34Hmtw9P9uyW5MGXVTT1vtd+/3b3evo/PTx/NaztN+y/an0rylqq6JMmnk/xYRtNIf1BVr0xyX5IfnvIxAFhccgKASeQEQA9M1Rxqrf1tku3GOF8yzf0CsBzkBACTyAmAfpjfLHzPDGlRNwAu3u23354kOXr0aJLHn9b7+c9//qHVstNI/UEszjkvWw/P2E73vYM6hGMW93vQ/3545JFHvvF1t4h3d+jC1sPItp5qeK+HwOx2CODW+4Ohuv7665NcONy4O3FBd6jnJz/5yUOvaevre96HOO3FxZ5koO+Hlc17W288scPWHNjJxeZfH7c7F2+aBakBAAAAWHAmh6BHJi2wB8zXZZddliR58MEHkyTHjx9PkjzwwLYn0ZnKbnsZt75XdHsCu6mmZZgc2rincyfTLrC81/ufZLfnatZ7jLvnuvu88bnuphR2yo6dapl2oel57DE+yMc86MWE7WFfXt1r8PTp09t+nqWLfV12f9fdVFO3gH0f7fU1stepxr47qKnX7vPG9/7u+d8pJ7w/7c0iTODth8khAAAAgAGrPnQHq2ruRVhziD4wOcR2WmuD/0OQE5N1j93tyVqWvak76aa2uqmZvUwZLbpuPamtUy0bT1HcfX3//fcfcnUXb+vrpQ//Hu306TTkeyUn+pETX//615M8fv2v7nM3WTRPyzrxsFX3HHQ5sey/b7LzunLdtFByIT+7dbDYn2XNCZNDAAAAAANmcmisz3uEGQ6TQ2zHHmE5wWbHjh1LcmFv6JkzZ5Is1h68vbr66quTJFdddVWS5LOf/WySC7/rok6JbZ2EWoZ1suZJTvQjJzaePTB5/HtSt3YdB69bg697j+nzGkvTevrTn54kufLKK5NcOCveQa/Ld9C6nOgsat71hckhAAAAACZytjIAYKFs3Qu6qHtFJ9m6dkQ3WbOI6xxMsiy/ByQX1rnZOrFhsvTwDWHbb82JbrKmz2u67cei179ITA4BAAAADJjJIeiRZd67Aezfsu0FnFa3bke3hsRhTdN0UwEbn4/HHnts022mfa66PcHd5+7+uj3C3fWLfoa2g1pL4mL+FpZtCov+6P62DvPfdXJis0svvTTJ49+jD+txN54h7PTp05tuM+1ztfXvq/v57nddlpzofo/u95vVGefkxM5MDgEAAAAMmMkhAOi5IaydcDEefPDBJIe/R6+bdpm0N3a/tXS/S/cY3R7S7nN3RrbD3gt+UGZ95qD9vDaGsieYw7N1DZjDJCc266Z1Dnt7dO/hk97j9vve0/19dVNJ3cTl1smhRZ8Y6sz69/Da2J3JIQAAAIABMzkEADNwGFMsJh02O+ztcRiP1z1Gt4f4oYceSjK7NXmA+ZETh++wt8dBTI9tvc9uorT7e3rggQc2Xc/sDO31ZHIIAAAAYMBMDgHADAxt79IQHT16NEmytrY28/veevaZbg9wnyeGhnYWF5iW18ryO3bsWJLZrqu2daJ0a070eWKoTzlhXa7dmRwCAAAAGLDqQxevquZexG7bQYcRmJfW2uDfgOTE5sfoQ3YzWztNDnFxLr300iTJo48+OudKDpeckBNbH0NOLI+tE0Pdcysn9qc729tBTAH32V5ywuQQAAAAwIBZcwgAFoQ9wcure249x/tz/PjxeZcAveA9ZHl1k0Ke4/3p1g1kZyaHAAAAAAbM5BAAAAut26N+5MiROVcCMFvWFpotawnvzOQQAAAAwICZHAIAltrWvYTWa1g+zvIGTENODIec2JnJIQAAAIABMzkEABeh25vY7WW0d3FxeK6WV7cn+OzZs3OuBB7Pe0//eY6WX5cT58+fn3Ml/WVyCAAAAGDATA4BwB6cPn1602V7GRePaa/lZWKIPnj00UeTPH4NLGdHWhxyYnmZGNqdySEAAACAATM5BAB7sLIy2p+ytraWJFldXd10Pf1njzBwkHbKCZND/ec5ApNDAAAAAIM21eRQVf1skn+XpCX5eJIfS/K0JG9L8qQkH07yb1prvT8QXLcYYPaWKSde/epXJ0nOnTuXJLn66quTJP/wD/8wt5rYm25SyMTQyNZpt25dFJiHZcqJ9773vUkurIF1xRVXJElOnTo1r5LYI/mwWZcT8nNY9j05VFXXJPnpJM9trT07yZEkNyf55SRvaK1dn+Qfk7xyFoUCsFjkBACTyAmA/pj2sLLVJMerajXJZUlOJXlxkrePv39HkpdP+RgALK6lyYkzZ87kzJkzOX36dE6fPp1Tp07l1KlTOX78eI4fPz7v8mDP1tfXN33AnC1NTjz66KObPh5++OE8/PDDOXbsWI4dOzbv8mDPWmubPhiGfTeHWmufS/KrSe7L6E38oSQfSvJga21tfLOTSa6ZtkgAFo+cAGASOQHQH9McVnZVkpuSXJfk6UkuT/LSbW66bauxqm6tqg9W1Qf3WwMA/SUnAJhETgD0xzQLUn9vks+01r6UJFX1R0m+K8mVVbU67vafSPL57X64tXZbktvGP2tWDWD5LFVO3HHHHfMuAWDZLFVO3HTTTfMuAWbCoWTDNM2aQ/cleUFVXVajU329JMknkrwnyQ+Nb3NLkndMVyIAC0pOADCJnADoiZqmK1hVv5DkXyVZS/KRjE5DeU0unHryI0n+dWvtsV3uR2sSYAettZp3DfslJwAOnpyQEwCT7CUnpmoOzYo3c4CdLfI/+mdFTgDsTE7ICYBJ9pIT057KHgAAAIAFpjkEAAAAMGCaQwAAAAADpjkEAAAAMGCaQwAAAAADpjkEAAAAMGCaQwAAAAADpjkEAAAAMGCaQwAAAAADpjkEAAAAMGCaQwAAAAADpjkEAAAAMGCaQwAAAAADpjkEAAAAMGCaQwAAAAADpjkEAAAAMGCaQwAAAAADpjkEAAAAMGCaQwAAAAADpjkEAAAAMGCaQwAAAAADpjkEAAAAMGCaQwAAAAADpjkEAAAAMGCaQwAAAAADpjkEAAAAMGCaQwAAAAADpjkEAAAAMGCaQwAAAAADpjkEAAAAMGCaQwAAAAADpjkEAAAAMGC7Noeq6s1V9UBV3bXhuidV1buq6lPjz1eNr6+q+s2qureqPlZVzznI4gGYPzkBwCRyAqD/9jI5dHuSG7Zc99ok726tXZ/k3ePLSfLSJNePP25N8sbZlAlAj90eOQHAzm6PnADotV2bQ6219yX56parb0pyx/jrO5K8fMP1v9dG3p/kyqp62qyKBaB/5AQAk8gJgP5b3efPPbW1dipJWmunquop4+uvSXL/htudHF93av8lArCA5ASwFFZWRvtSq2rb758/f/4wy1kmcgJYCpdcckmSC3mxvr6+6ftnz5499Jr2Y7/NoZ1sl5pt2xtW3ZrRqCgAwyEnAJhETgDMwX7PVvbFbrxz/PmB8fUnk1y74XYnknx+uztord3WWntua+25+6wBgP6SE8BSaK1t+th6PfsmJ4ClsL6+vumjqlJVC5cT+20O3ZnklvHXtyR5x4brf3R8loEXJHmoGxcFYFDkBACTyAmAHqndOllV9dYkL0ry5CRfTPLzSf4kyR8k+aYk9yX54dbaV2t0MPZ/y+hsBGeS/Fhr7YO7FlG1OO00gEPWWtt+oYuekBPAMnve856XJLn88suTXFhLoltr6K/+6q/mU9gGckJOAPPzEz/xE0mSJz7xiUmS48ePJ7mw1tAv/uIvzqewDfaSE7uuOdRa+5EdvvWSbW7bkvzk7qUBsCzkBACTyAmA/pv1gtQAALA0uj2/R48eTZKsro7++XzmzJm51QRAf3T50J3V8pFHHkmSPPzww3OraT/2u+YQAAAAAEtg1zWHDqUIxwgDc7KHddcOqZKd9X0ticMgJ4B5efDBB5MkR44c2XT92tpakuSqq6469Jq2khNyApifc+fOJbnw/4buc7dGXTdZNE97yQmTQwAAAAADZs0hAADYQTdhurKyeZ9qt/YQACSPP+KgD0cgXAyTQwAAAAADZs0hYNCsObQY5AQwL92aETvZOlE0D3JCTgDzsyz/n5h/mgEAAAAwN5pDAAAAAAOmOQQAAAAwYE6zAAAAO+jDWhEAcNBMDgEAAAAMmOYQAAAAwIBpDgEAAAAMmOYQAAAAwIBpDgEAAAAMmOYQAAAAwIBpDgEAAAAMmOYQAAAAwIBpDgEAAAAMmOYQAAAAwIBpDgEAAAAMmOYQAAAAwIBpDgEAAAAMmOYQAAAAwIBpDgEAAAAMmOYQAAAAwIBpDgEAAAAMmOYQAAAAwICtzrsAAABYNK21eZcAADNjcggAAABgwEwOAQDADs6ePTvvEgDosXPnzm26XFVzqmQ6u04OVdWbq+qBqrprw3W/UlWfrKqPVdUfV9WVG773uqq6t6ruqarvP6jCAegHOQHAJHICoP/2cljZ7Ulu2HLdu5I8u7X27Un+PsnrkqSqvjXJzUm+bfwzv11VR2ZWLQB9dHvkBLDkWmub1hnaepmJbo+cAJZclwtbPxbFrs2h1tr7knx1y3V/0VpbG198f5IT469vSvK21tpjrbXPJLk3yfNmWC8APSMnAJhETgD03yzWHPrxJL8//vqajN7cOyfH1wH00qIeE7xg5ASwsF7xilckSZ7whCdsuv706dPzKGdZyQlgYd14441JkmPHjiVJjh49miR5+OGH51bTfkzVHKqq1ydZS/KW7qptbrbtHFVV3Zrk1mkeH4B+kxMATCInAPph382hqrolycuSvKRdOJDuZJJrN9zsRJLPb/fzrbXbktw2vq/FORAPgD2RE8AyOHJktNzN2troCKgvf/nL8yxnqcgJYBl0ObG6OmqvfOELX9h0eVHsZUHqx6mqG5K8JsmNrbUzG751Z5Kbq+pYVV2X5Pokfz19mQAsEjkBwCRyAqBfarfVs6vqrUlelOTJSb6Y5OczOpvAsSRfGd/s/a21fz++/eszOm54LcnPtNb+dNcidPoBdtRa6/XCSHICWGZXXHFFkuTcuXNJkksuuSTJhUmiPqwpISfkBDA/l112WZILudCtOXT+/PkkyaOPPjqfwjbYS07s2hw6DN7MAXbW93/0HwY5AcyL5tBikBPAvCxLc2ixDoIDAIBDtPUf+2fPnk1yYY0JAIatawJ1awx1OxNWVva1is/cLFa1AAAAAMyUw8oAes7hAnICYBI5IScAJtlLTpgcAgAAABgwzSEAAACAAdMcAgAAABgwzSEAAACAAdMcAgAAABgwzSEAAACAAdMcAgAAABgwzSEAAACAAdMcAgAAABgwzSEAAACAAdMcAgAAABgwzSEAAACAAdMcAgAAABgwzSEAAACAAdMcAgAAABgwzSEAAACAAdMcAgAAABgwzSEAAACAAdMcAgAAABgwzSEAAACAAdMcAgAAABgwzSEAAACAAdMcAgAAABgwzSEAAACAAdMcAgAAABgwzSEAAACAAdMcAgAAABgwzSEAAACAAdMcAgAAABiwXZtDVfXmqnqgqu7a5nuvrqpWVU8eX66q+s2qureqPlZVzzmIogHoDzkBwCRyAqD/9jI5dHuSG7ZeWVXXJvm+JPdtuPqlSa4ff9ya5I3TlwhAz90eOQHAzm6PnADotV2bQ6219yX56jbfekOSn0vSNlx3U5LfayPvT3JlVT1tJpUC0EtyAoBJ5ARA/+1rzaGqujHJ51prH93yrWuS3L/h8snxdQAMiJwAYBI5AdAvqxf7A1V1WZLXJ/mX2317m+vaNtelqm7NaFQUgCUiJwCYRE4A9M9FN4eSfHOS65J8tKqS5ESSD1fV8zLq7F+74bYnknx+uztprd2W5LYkqapt3/ABWEhyAoBJ5ARAz1z0YWWttY+31p7SWntGa+0ZGb2BP6e19oUkdyb50fFZBl6Q5KHW2qnZlgxAn8kJACaREwD9s5dT2b81yf9N8qyqOllVr5xw8/+d5NNJ7k3yP5L8h5lUCUBvyQkAJpETAP1Xrc1/AtMYKMDOWmvbrb8wKHICYGdyQk4ATLKXnNjX2coAAAAAWA6aQwAAAAADpjkEAAAAMGCaQwAAAAADpjkEAAAAMGCaQwAAAAADpjkEAAAAMGCaQwAAAAADpjkEAAAAMGCaQwAAAAADpjkEAAAAMGCaQwAAAAADpjkEAAAAMGCaQwAAAAADtjrvAgAA4CCtrFzYH9pa2/QZAOSEySEAAACAQTM5BADAUltfX593CQD02NCmhLZjcggAAABgwDSHAAAAAAbMYWUAAADAYDmszOQQAAAAwKBpDgEAAAAMmOYQAAAAwIBZcwgAAABgH6pq4vcXZT0jk0MAAAAAA2ZyCADggHV7FRdl7yH9srIy2p+7vr4+50qAWevyofvsdb64uozfbZLoIMwiJ0wOAQAAAAyYySEAgANmYmhx9WEtCZMEsPy8zhdXN7WzUx4sSk6YHAIAAAAYMJNDAAAHxFpDy8tzCsySvFhcO601tGjTYCaHAAAAAAbM5BAAwAGxB3jxdWtJdLrntLv+/PnzB/7Yi7b3Gdg7ObH4jhw5suly95yuro7aLefOnTuwx55lTpgcAgAAABgwk0MAALBHu52VZpZMDAH039Y1hxY1J0wOAQAAAAxYXyaHvpzk4fHnPnty1DgLapwNNc5G32v8p/MuoCfkxOyocTbUOBu9r/H8+fN9r1FOjMiJ2VHjbKhxNnpf49raWt9r3FNOVF8WwKqqD7bWnjvvOiZR42yocTbUOBuLUCMji/BcqXE21DgbapyNRaiRkUV4rtQ4G2qcDTXOxiLUuBcOKwMAAAAYMM0hAAAAgAHrU3PotnkXsAdqnA01zoYaZ2MRamRkEZ4rNc6GGmdDjbOxCDUysggUCElwAAAGSklEQVTPlRpnQ42zocbZWIQad9WbNYcAAAAAOHx9mhwCAAAA4JD1ojlUVTdU1T1VdW9VvXbe9SRJVV1bVe+pqrur6u+q6lXj659UVe+qqk+NP1815zqPVNVHquqd48vXVdUHxvX9flVdMs/6xjVdWVVvr6pPjrfnC3u4HX92/DzfVVVvrapL570tq+rNVfVAVd214bptt1uN/Ob4NfSxqnrOHGv8lfFz/bGq+uOqunLD9143rvGeqvr+edW44XuvrqpWVU8eX57LdmR3cmKqOuXEbGqUE7OrUU4wUzJi6lrlxPT19S4jxnXJiQOqccP3liYn5t4cqqojSX4ryUuTfGuSH6mqb51vVUmStST/ubX2z5K8IMlPjut6bZJ3t9auT/Lu8eV5elWSuzdc/uUkbxjX949JXjmXqjb7jSR/1lr7liTfkVG9vdmOVXVNkp9O8tzW2rOTHElyc+a/LW9PcsOW63babi9Ncv3449Ykb5xjje9K8uzW2rcn+fskr0uS8evn5iTfNv6Z3x6//udRY6rq2iTfl+S+DVfPazsygZyYmpyYkpyYeY1ygpmRETMhJ6bQ44xI5MRB1rh0OTH35lCS5yW5t7X26dba2SRvS3LTnGtKa+1Ua+3D46+/ntEb0DUZ1XbH+GZ3JHn5fCpMqupEkh9M8qbx5Ury4iRvH99krvUlSVVdkeRfJPmdJGmtnW2tPZgebcex1STHq2o1yWVJTmXO27K19r4kX91y9U7b7aYkv9dG3p/kyqp62jxqbK39RWttbXzx/UlObKjxba21x1prn0lyb0av/0OvcewNSX4uycaF1+ayHdmVnNgnOTFTcmJGNcoJZkxGTEFOzEzvMiKREwdZ49hS5UQfmkPXJLl/w+WT4+t6o6qekeQ7k3wgyVNba6eS0Zt+kqfMr7L8ekZ/jOvjy1cneXDDC6kP2/KZSb6U5HfH46pvqqrL06Pt2Fr7XJJfzajjeyrJQ0k+lP5ty2Tn7dbX19GPJ/nT8de9qbGqbkzyudbaR7d8qzc1sknvnxc5MRU5MVtyYgbkxELp/XPS44xI5MTUFiwjEjkxE8uYE31oDtU21/XmFGpV9YQkf5jkZ1prX5t3PZ2qelmSB1prH9p49TY3nfe2XE3ynCRvbK19Z5KH04/x2W8YH2d7U5Lrkjw9yeUZjQNuNe9tOUnvnvuqen1GI9Vv6a7a5maHXmNVXZbk9Un+63bf3ua6Pj/vQ9Hr50VOTE1OHI7ePfdyghnp9XPS14xI5MSsLElGJD187uXE4epDc+hkkms3XD6R5PNzqmWTqjqa0Zv5W1prfzS++ovdWNj48wNzKu+7k9xYVZ/NaHz2xRl1/q8cjzMm/diWJ5OcbK19YHz57Rm9ufdlOybJ9yb5TGvtS621c0n+KMl3pX/bMtl5u/XqdVRVtyR5WZJXtNa6N8O+1PjNGYX3R8evnxNJPlxV/yT9qZHNevu8yImZkBOzJSemJycWS2+fk55nRCInZmWRMiKRE7OwlDnRh+bQ3yS5vkaruV+S0QJTd865pu54299Jcndr7dc2fOvOJLeMv74lyTsOu7Ykaa29rrV2orX2jIy22V+21l6R5D1Jfmje9XVaa19Icn9VPWt81UuSfCI92Y5j9yV5QVVdNn7euxp7tS3Hdtpudyb50fHq+C9I8lA3LnrYquqGJK9JcmNr7cyGb92Z5OaqOlZV12W0SNtfH3Z9rbWPt9ae0lp7xvj1czLJc8Z/q73ZjmwiJ/ZBTsyUnJghOcGMyYh9khMzs0gZkciJqS1tTrTW5v6R5AcyWoX8/yV5/bzrGdf0PRmNf30syd+OP34go+Nw353kU+PPT+pBrS9K8s7x18/M6AVyb5L/leRYD+r750k+ON6Wf5Lkqr5txyS/kOSTSe5K8j+THJv3tkzy1oyOWz6X0RvOK3fabhmNL/7W+DX08YzOljCvGu/N6Djb7nXz3zfc/vXjGu9J8tJ51bjl+59N8uR5bkcfe3oe5cR0tcqJ6WuUE7OrUU74mPVzKCOmr1dOTFdf7zJiXJecOKAat3x/KXKixr8AAAAAAAPUh8PKAAAAAJgTzSEAAACAAdMcAgAAABgwzSEAAACAAdMcAgAAABgwzSEAAACAAdMcAgAAABgwzSEAAACAAfv/i1Govi8ETPkAAAAASUVORK5CYII=\n",
      "text/plain": [
       "<Figure size 1440x720 with 3 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "image, label = iter(Pong_train).next()\n",
    "with torch.no_grad():\n",
    "    mu_z1, log_std_z1, mu_x1, log_std_x1, z1, x1 = Pong_VAE1(image[0,:,:].view(-1), noise=False)\n",
    "    mu_z2, log_std_z2, mu_x2, log_std_x2, z2, x2 = Pong_VAE1(image[0,:,:].view(-1), noise=True)\n",
    "x1 = x1.view(160,160)\n",
    "x2 = x2.view(160,160)\n",
    "plt.figure(figsize=(20,10))\n",
    "plt.subplot(1,3,1)\n",
    "plot_image(tensor=image[0,:,:])\n",
    "plt.subplot(1,3,2)\n",
    "plot_image(x1)\n",
    "plt.subplot(1,3,3)\n",
    "plot_image(x2)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In the above example we can see that VAE fails to recover the position with a ball: It only have a very dark mark on that position instead of a while point. But in the following example you can see that the ball is recovered."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAABIcAAAFrCAYAAABLz+nFAAAABHNCSVQICAgIfAhkiAAAAAlwSFlzAAALEgAACxIB0t1+/AAAADl0RVh0U29mdHdhcmUAbWF0cGxvdGxpYiB2ZXJzaW9uIDMuMC4yLCBodHRwOi8vbWF0cGxvdGxpYi5vcmcvOIA7rQAAIABJREFUeJzt3X+MrXddJ/D3p3P7G7QFUhZ7Ky3aoEDUVsLirw0BXYsQyh+arXFjo2yazboRXY3CkqzxT6MRNSqbLmDRGJDFHzQEf5CKYf9YkCJLoS1IF0y5cKVsKpW0cnvn3u/+MefpnTvMnJnOOTPPc+b7eiWTmfOcMzOfe545533u5/mc71OttQAAAADQpwvGLgAAAACA8WgOAQAAAHRMcwgAAACgY5pDAAAAAB3THAIAAADomOYQAAAAQMc0hwAAAAA6dmDNoaq6qao+VVUPVNXrDur3ALB6ZAQA88gJgMNVrbXl/9CqtSR/n+QHkpxI8uEkP9pau2/pvwyAlSIjAJhHTgAcvoOaHHpRkgdaa59prT2e5B1Jbj6g3wXAapERAMwjJwAO2bED+rlXJ/ncpssnkvzrnW5cVcsfXwLYg+/8zu+ce/1HPvKRQ6pkZ621GruGJXtSGZFMIycuuGD+8ZSzZ88eUiXAYZITo1jJnKja2A3DOzOGy4ODeMcGML6jkhMH1Rza7hef92xYVbclue2Afj/Antx9991zr9/6wo6l2DUjkunnxE4v/oGjRU6MYiVzYmvzRzMI+vDhD3947vW7HWCcioNqDp1Ics2my8eTfGHzDVprtye5PZlGpx+AQ7NrRiRyAqBjcgLgkB1UC+vDSa6vquuq6qIktyS584B+FwCrZSUzorV23kdVpaqeuAzA0qxkTgB9Gl4T7vSxKg5kcqi1tl5V/znJXyZZS/LW1tq9B/G7AFgtMgKAeeQEwOE7kFPZP+kijIECI9ntOXAK3f4juNDokyYngLHIidUgJ4CxHJWcWI2VkQAAAAA4EJpDAAAAAB3THAIAAADomOYQAAAAQMc0hwAAAAA6pjkEAAAA0DHNIQAAAICOaQ4BAAAAdExzCAAAAKBjmkMAAAAAHdMcAgAAAOiY5hAAAABAxzSHAAAAADqmOQQAAADQsWNjFwAAAFPVWhu7BAA4cCaHAAAAADqmOQQAAADQMc0hAAAAgI5ZcwgAAHZw9uzZ8y4PaxBV1RjlAMCBMDkEAAAA0DGTQwAAsIthUmiYHHIWMwCOEpNDAAAAAB0zOQQAADu45557kiSPPvpokuTCCy9Mklx++eWj1QTAdDz22GNJkjNnziQ5N2l67NhqtVtMDgEAAAB0bLVaWQAAcIhOnTqVJLn00kuTJBdffHGS5PTp06PVBMB0DGe1HCaG1tbWkiTr6+uj1bQfJocAAAAAOmZyCOja0OEHgO183/d933mXr7zyyiTJs5/97DHKAWBivvmbvznJuTWHbrzxxiTJDTfcMFpN+2FyCAAAAKBjJocAAGCPLrhg49hqa23kSgCYgmGNoWHtocsuuyzJuTXrVoXJIQAAAICO1RSOelTV+EUATFRrrfuFkeQEMLZjx46d93lYW2IKZy2TE3ICGM+whukwMXThhRcmOZcTX/nKV8YpbJO95ITJIQAAAICO7bs5VFXXVNX7q+r+qrq3ql472/60qnpfVX169vnK5ZULwKqQE8BRdObMmSeOBrMYOQEcBa218z5WNScWmRxaT/JzrbVvTfLiJD9VVc9L8rokd7XWrk9y1+wyAP2REwDMIycAJmJpaw5V1buT/Pbs4yWttZNV9awkf9Nae+4u3+s9wgA7OCprScgJYJVtPUvZFNbtHMgJOQGMZ1hzaMiJwZSmhw5tzaGqujbJDUk+lOSZrbWTswJOJrlqGb8DgNUlJwCYR04AjOvYoj+gqp6S5I+T/Exr7Z+Hrtkevu+2JLct+vsBmDY5ARwFU5oUOmrkBHAUrHpOLPS2sqq6MMl7kvxla+3XZ9s+FWOgAEuzym8XkBPAUbG1YTGl/wTICTkBjGfIh605cfbs2THK2daBvq2sNv7lb0ly//BEPnNnkltnX9+a5N37/R0ArC45ARwlW89Gw+LkBHAUDLlw9uzZ8z5Wzb4nh6rqe5P8ryQfTzL8y/9rNt4n/M4k35jkwSQ/0lp7eJefJWEBdrCqR4TlBMDhkBNyAmCeveTE0s5WtghP5gA7W9UX/cskJwB2JifkBMA8h3a2MgAAAABWk+YQAAAAQMc0hwAAAAA6pjkEAAAA0DHNIQAAAICOaQ4BAAAAdExzCAAAAKBjmkMAAAAAHdMcAgAAAOiY5hAAAABAxzSHAAAAADqmOQQAAADQMc0hAAAAgI5pDgEAAAB0THMIAAAAoGOaQwAAAAAd0xwCAAAA6JjmEAAAAEDHNIcAAAAAOqY5BAAAANAxzSEAAACAjmkOAQAAAHRMcwgAAACgY5pDAAAAAB3THAIAAADomOYQAAAAQMc0hwAAAAA6pjkEAAAA0DHNIQAAAICOaQ4BAAAAdExzCAAAAKBjmkMAAAAAHVu4OVRVa1X10ap6z+zydVX1oar6dFX9UVVdtHiZAKwqOQHAPHICYHzLmBx6bZL7N13+lSRvbK1dn+SfkrxmCb8DgNUlJwCYR04AjGyh5lBVHU/yiiRvnl2uJC9N8q7ZTd6W5NWL/A4AVpecAGAeOQEwDYtODv1Gkl9IcnZ2+elJvtxaW59dPpHk6gV/BwCrS04AMI+cAJiAfTeHquqVSR5qrX1k8+Ztbtp2+P7bquruqrp7vzUAMF1yAoB55ATAdBxb4Hu/J8mrquqHklyS5Ouy0fm/oqqOzbr9x5N8Ybtvbq3dnuT2JKmqbZ/wAVhpcgKAeeQEwETse3Kotfb61trx1tq1SW5J8tettR9L8v4kPzy72a1J3r1wlQCsHDkBwDxyAmA6lnG2sq1+Mcl/qaoHsvGe4bccwO8AYHXJCQDmkRMAh6xaG38C0xgowM5aa9utv9AVOQGwMzkhJwDm2UtOHMTkEAAAAAArQnMIAAAAoGOaQwAAAAAd0xwCAAAA6JjmEAAAAEDHNIcAAAAAOqY5BAAAANAxzSEAAACAjmkOAQAAAHRMcwgAAACgY5pDAAAAAB3THAIAAADomOYQAAAAQMc0hwAAAAA6pjkEAAAA0DHNIQAAAICOaQ4BAAAAdExzCAAAAKBjmkMAAAAAHdMcAgAAAOiY5hAAAABAxzSHAAAAADqmOQQAAADQMc0hAAAAgI5pDgEAAAB0THMIAAAAoGOaQwAAAAAd0xwCAAAA6JjmEAAAAEDHNIcAAAAAOqY5BAAAANAxzSEAAACAji3UHKqqK6rqXVX1yaq6v6q+q6qeVlXvq6pPzz5fuaxiAVgtcgKAeeQEwDQsOjn0m0n+orX2LUm+Pcn9SV6X5K7W2vVJ7ppdBqBPcgKAeeQEwARUa21/31j1dUk+luQ5bdMPqapPJXlJa+1kVT0ryd+01p67y8/aXxEAHWit1dg17IecADgcckJOAMyzl5xYZHLoOUm+lOT3quqjVfXmqro8yTNbaydnBZxMctUCvwOA1SUnAJhHTgBMxCLNoWNJbkzyptbaDUkezZMY+ayq26rq7qq6e4EaAJguOQHAPHICYCIWaQ6dSHKitfah2eV3ZePJ/Yuz8c/MPj+03Te31m5vrb2wtfbCBWoAYLrkBADzyAmAidh3c6i19o9JPldVw/t/X5bkviR3Jrl1tu3WJO9eqEIAVpKcAGAeOQEwHftekDpJquo7krw5yUVJPpPkJ7LRcHpnkm9M8mCSH2mtPbzLz7GAHMAOVnWh0UROABwGOSEnAObZS04s1BxaFk/mADtb5Rf9yyInAHYmJ+QEwDwHfbYyAAAAAFac5hAAAABAxzSHAAAAgG5UVaq6f0fueTSHAAAAADp2bOwCAAAAAA7LFE7MNTUmhwAAAAA6pjkEAAAA0DHNIQAAAICOaQ4BAAAAdExzCAAAAKBjmkMAAAAAHdMcAgAAAOiY5hAAAABAxzSHAAAAADqmOQQAAADQMc0hAAAAgI5pDgEAAAB0THMIAAAAoGOaQwAAAAAd0xwCAAAA6JjmEAAAAEDHNIcAAAAAOqY5BAAAANAxzSEAAACAjmkOAQAAAHRMcwgAAACgY5pDAAAAAB3THAIAAADomOYQAAAAQMc0hwAAAAA6pjkEAAAA0DHNIQAAAICOLdQcqqqfrap7q+oTVfX2qrqkqq6rqg9V1aer6o+q6qJlFQvAapETAMwjJ6A/VfXEB9Ox7+ZQVV2d5KeTvLC19oIka0luSfIrSd7YWrs+yT8lec0yCgVgtcgJAOaREwDTsejbyo4lubSqjiW5LMnJJC9N8q7Z9W9L8uoFfwcAq0tOADCPnFgCUxisktbaEx/+dqdj382h1trnk/xakgez8ST+SJKPJPlya219drMTSa5etEgAVo+cAGAeOQEwHYu8rezKJDcnuS7JNyS5PMnLt7lp2+H7b6uqu6vq7v3WAMB0yQkA5pETANNxbIHv/f4kn22tfSlJqupPknx3kiuq6tis2388yRe2++bW2u1Jbp9977ZP+ACsNDkBwDxyYkla6/qfzwrztzsdi6w59GCSF1fVZbXxJsGXJbkvyfuT/PDsNrcmefdiJQKwouQEAPPICYCJqEU6dVX1y0n+XZL1JB9N8h+y8Z7gdyR52mzbv2+tndrl52gXAuygtbayq/TJCYDtra2tPfH1mTNnFvpZckJOAEfPBRecm+U5e/bsQj9rLzmxUHNoWTyZA+xslV/0L4ucAI4azaHlkhPAUXPYzaFF1hwCAAD2YdEX+rAXwynCpzAQADw5h/24XWTNIQAAAABWnMkhAFgxjgRPl33Dboa3k5kc4iANz0XA6hkrJ0wOAQAAAHTM5BAArBhTKdM17BsTROzE3waHyd8ZrC5rDgEAAABwaEwOAcCKGU5tas2S6RqO9l166aVJzu2rU6dOjVYT4xoet4ueth72YuuaQyaIpuuiiy5Kcm4fnT59esxyGNHYa9KZHAIAAADomMkhANiDSy655LzLW6d3vvrVrx5aLSaGVsfwdzIcxT92bOOl1/r6+mg1MQ6TQ0ffsI8HW9eXOsznbjmxOra+nhgmiR5//PHRamIcY0+GmxwCAAAA6JjJIQDYgymfheog1pa47LLLkpybcHEEc3+2Tgw4mt8v02KM6SByYpioHabhrJWzP8O+GD6bLhzH5sm/sbJ67MeQySEAAACAjpkcmtmte7612w5AX06cOJHk3JGlYQpgWEPm6U9/+jiFbXIQa5psXUODJ2fYF15HwNH33ve+N8m5sxQ++uijSc5N2Lz0pS8dp7BNhsxa5oSCnFjMMKUiJ8Y1pYnwsXgkAwAAAHTM5BAcoK0daEcEYHUNR/aGo65TOlJ6EGfC2boGwiqawvpQp06dSpKsra2dV9Pgi1/8YpJzZ6e58sorD7E6DsKwr60b0p9hbbanPvWpSc79DQx/E1OwzL/LMc7CtmxTyIlhimunv5N77703SXL55ZcnSa699tpDqYuDM9WcmM4rWwAAAAAOnckhANiDq666auwSdnQQU4kmHZdrt6ODw4TRFI5isxj7rl+veMUrknztGjJTeD49yBqm8O87CnbKiWFS+V/+5V+SyImD4v40OQQAAADQNZNDALDiDuJo12OPPbb0n8nXGtYdGNaycuRy9Q1TI8PR/qOwfhd7M+zjKe77g1gXaJh45GANfz9PecpTzrvM6hqmxLauXzn2+l0mhwAAAAA6ZnIIAGBk6+vrY5fAAbnwwguTnDuTFUfXss5KtnXyaOvX9OWSSy5Jknz1q19NYs2ho2g4Y+mwj8dicggAAACgYzWFjmNVjV7EbveDVfiBsbTWun8CkhMcdcOkwNb1Boa1iAYmjFbPMDl0+vTpA/sdcuJo5sTm20/h/2yMa6ecmNq6NTx5W6fDDsJecsLkEAAAAEDHrDkEADCynY70mhZYffYh++Vvh83kxNE1lX1ocggAAACgYyaHAAAmajiauHVNCVaHdaKAw+AsZuMZMvqpT31qknNrzJ06dSpJcubMmbnfP9xubF5pAAAAAHTM5BAAwEStra0l2f2oIwB9MjE0HUNWD9m9ameyNTkEAAAA0DHNIQAAAICOeVsZAMBEDYtaMr7hbQLDwqP2DTAFO53insNz+eWXJ0me//znJzmXD5/97GeTJA8//PA4hT1Ju04OVdVbq+qhqvrEpm1Pq6r3VdWnZ5+vnG2vqvqtqnqgqu6pqhsPsngAxicnAJhHTgBM317eVnZHkpu2bHtdkrtaa9cnuWt2OUlenuT62cdtSd60nDIBmLA7Iid2VFXnfXA02b9H37CPW2sWfn3y7oic2NEFF1xw3gdHk5w4us6ePZuzZ88+kQ+nTp3KqVOnVi4vdn32aa19IMnWOaibk7xt9vXbkrx60/bfbxs+mOSKqnrWsooFYHrkBADzyAmA6dvvmkPPbK2dTJLW2smqumq2/eokn9t0uxOzbSf3XyIAK0hOzKzSESNgZ+vr62OXcNTIiRlrxsBqG17r3XfffedtP3ZstZZ4Xna1283JbfuquKpuy8aoKAD9kBMAzCMnAEaw3+bQF6vqWbMu/7OSPDTbfiLJNZtudzzJF7b7Aa2125PcniRV5bAqwNFy5HLCBBDz+PuAJ+3I5cSZM2fOuzxMBFlnhkROHGXDhNDwmB8uP/7446PVtB/7XfHsziS3zr6+Ncm7N23/8dlZBl6c5JFhXBSArsgJAOaREwATUrt1MKvq7UlekuQZSb6Y5JeS/FmSdyb5xiQPJvmR1trDtdEW/+1snI3gsSQ/0Vq7e9ciJtDp38P9cEiVAJyvtTbpJ6BecmK3NSGcYQYYi5yYRk5snRIY/n8x/D/ioosuOvSagIM3PMbX1taSfO3U4NapwjHsJSd2bQ4dhik8mWsOAVM19Rf9h2EKOaE5BEyVnJhGTmgOQZ+OSnNotZbPBoCRnDp1Ksm5F/vDC4ApHGRheoYXhP4+oB+PPPJIkuT06dNJvnYdEths6/CBvFhdW/flqr4GcJgTAAAAoGMmhwBgD7YeBRqOBK/aUSEADsb6+nqSc3kx5MQU3lLC9Gz9O2F1Dftw1felySEAAACAjlmQGmDiLDQ6jZz4+q//+iTnJoWGhUWHo0QPP/zwOIUxaau67gCrRU5MIycuvfTSJF+7OO0wOfTYY4+NUxiTJic4DHvJCZNDAAAAAB2z5hAA7MFwRHg4VfFwJHhYYwKAvl144YVJzuXDUVmHhIMxTAyZHGIqTA4BAAAAdMyaQwATZy0JOQEwj5yQEwDzWHMIAAAAgLk0hwAAAAA6pjkEAAAA0DHNIQAAAICOaQ4BAAAAdExzCAAAAKBjmkMAAAAAHdMcAgAAAOiY5hAAAABAxzSHAAAAADqmOQQAAADQsWNjFwAAcBCqKknSWhu5EgCm6IILNmYlhpyQF/TM5BAAAABAxzSHAAAAADqmOQQAAADQMWsOAQBH0rDm0PD57NmzY5YDwMQM+TCsPbS+vj5mOTAqk0MAAAAAHdMcAgAAAOiYt5UBAEeSUxIDMM/wtjJ5ASaHAAAAALpmcggAOJIcCQZgHgtQwzkmhwAAAAA6pjkEAAAA0LFdm0NV9daqeqiqPrFp269W1Ser6p6q+tOqumLTda+vqgeq6lNV9YMHVTgA0yAnAJhHTgBM314mh+5IctOWbe9L8oLW2rcl+fskr0+SqnpekluSPH/2Pb9bVWtLqxaAKbojcgKAnd0ROQEwabs2h1prH0jy8JZtf9VaG1bv+mCS47Ovb07yjtbaqdbaZ5M8kORFS6wXgImREwDMIycApm8Zaw79ZJI/n319dZLPbbruxGwbAP2SEwDMIycARrbQqeyr6g1J1pP84bBpm5ttex7ZqrotyW2L/H4Apk1OADCPnACYhn03h6rq1iSvTPKy1trwhH0iyTWbbnY8yRe2+/7W2u1Jbp/9rG2f8AFYXXICgHnkBMB07OttZVV1U5JfTPKq1tpjm666M8ktVXVxVV2X5Pokf7t4mQCsEjkBwDxyAmBadp0cqqq3J3lJkmdU1Ykkv5SNswlcnOR9VZUkH2yt/cfW2r1V9c4k92VjPPSnWmtnDqp4AMYnJwCYR04ATF+dm+AcsQhjoAA7aq1tt/5CV+QEwM7khJwAmGcvObGMs5UBAAAAsKI0hwAAAAA6pjkEAAAA0DHNIQAAAICOaQ4BAAAAdExzCAAAAKBjmkMAAAAAHdMcAgAAAOiY5hAAAABAxzSHAAAAADqmOQQAAADQMc0hAAAAgI5pDgEAB66qUlVjlwHARMkJGJfmEAAAAEDHjo1dAABw9LXWxi6BI2xtbS1JcsEFG8c9T58+PWY5AEzMkBPDdNr6+vqY5UySySEAAACAjmkOAQCw0oa1StbW1p44OgysltaaKVMOnJzYmeYQAAAAQMesOQQAwEozbQAAizE5BAAAANAxk0MAAKy0YXLI2WcA2M6QE2fOnBm5kukyOQQAAADQMc0hAAAAgI55WxkAACvt7Nmz530GgM3kxO5MDgEAAAB0THMIAAAAoGOaQwAAAAAd0xwCAAAA6JjmEAAAAEDHNIcAAAAAOqY5BAAAANAxzSEAAACAju3aHKqqt1bVQ1X1iW2u+/mqalX1jNnlqqrfqqoHquqeqrrxIIoGYDrkBADzyAmA6dvL5NAdSW7aurGqrknyA0ke3LT55Umun33cluRNi5cIwMTdETkBwM7uiJwAmLRdm0OttQ8keXibq96Y5BeStE3bbk7y+23DB5NcUVXPWkqlAEySnABgHjkBMH37WnOoql6V5POttY9tuerqJJ/bdPnEbBsAHZETAMwjJwCm5diT/YaquizJG5L82+2u3mZb22Zbquq2bIyKAnCEyAkA5pETANPzpJtDSb4pyXVJPlZVSXI8yd9V1Yuy0dm/ZtNtjyf5wnY/pLV2e5Lbk6Sqtn3CB2AlyQkA5pETABPzpN9W1lr7eGvtqtbata21a7PxBH5ja+0fk9yZ5MdnZxl4cZJHWmsnl1syAFMmJwCYR04ATM9eTmX/9iT/O8lzq+pEVb1mzs3fm+QzSR5I8j+S/KelVAnAZMkJAOaREwDTV62NP4FpDBRgZ6217dZf6IqcANiZnJATAPPsJSf2dbYyAAAAAI4GzSEAAACAjmkOAQAAMAlVldlZ7IBDpDkEAAAA0LFjYxcAAAAASTKFEyZBj0wOAQAAAHRMcwgAAACgY5pDAAAAAB3THAIAAGASnK0MxqE5BAAAANAxZysDAABgEpytDMZhcggAAACgY5pDAAAAAB3THAIAAADomOYQAAAAQMc0hwAAAAA6pjkEAAAA0DHNIQAAAICOaQ4BAAAAdExzCAAAAKBjmkMAAAAAHdMcAgAAAOiY5hAAAABAxzSHAAAAADqmOQQAAADQMc0hAAAAgI5pDgEAAAB0THMIAAAAoGOaQwAAAAAd0xwCAAAA6JjmEAAAAEDHNIcAAAAAOqY5BAAAANAxzSEAAACAjmkOAQAAAHTs2NgFzPy/JI/OPk/ZM6LGZVDjcqhxOaZe47PHLmAi5MTyqHE51LgcalycnNggJ5ZHjcuhxuVQ4+L2lBPVWjvoQvakqu5urb1w7DrmUeNyqHE51Lgcq1AjG1ZhX6lxOdS4HGpcjlWokQ2rsK/UuBxqXA41Lscq1LgX3lYGAAAA0DHNIQAAAICOTak5dPvYBeyBGpdDjcuhxuVYhRrZsAr7So3LocblUONyrEKNbFiFfaXG5VDjcqhxOVahxl1NZs0hAAAAAA7flCaHAAAAADhkk2gOVdVNVfWpqnqgql43dj1JUlXXVNX7q+r+qrq3ql472/60qnpfVX169vnKketcq6qPVtV7Zpevq6oPzer7o6q6aMz6ZjVdUVXvqqpPzu7P75rg/fizs/38iap6e1VdMvZ9WVVvraqHquoTm7Zte7/Vht+aPYbuqaobR6zxV2f7+p6q+tOqumLTda+f1fipqvrBsWrcdN3PV1WrqmfMLo9yP7I7ObFQnXJiOTXKieXVKCdYKhmxcK1yYvH6JpcRs7rkxAHVuOm6I5MTozeHqmotye8keXmS5yX50ap63rhVJUnWk/xca+1bk7w4yU/N6npdkrtaa9cnuWt2eUyvTXL/psu/kuSNs/r+KclrRqnqfL+Z5C9aa9+S5NuzUe9k7sequjrJTyd5YWvtBUnWktyS8e/LO5LctGXbTvfby5NcP/u4LcmbRqzxfUle0Fr7tiR/n+T1STJ7/NyS5Pmz7/nd2eN/jBpTVdck+YEkD27aPNb9yBxyYmFyYkFyYuk1ygmWRkYshZxYwIQzIpETB1njkcuJ0ZtDSV6U5IHW2mdaa48neUeSm0euKa21k621v5t9/ZVsPAFdnY3a3ja72duSvHqcCpOqOp7kFUnePLtcSV6a5F2zm4xaX5JU1dcl+TdJ3pIkrbXHW2tfzoTux5ljSS6tqmNJLktyMiPfl621DyR5eMvmne63m5P8ftvwwSRXVNWzxqixtfZXrbX12cUPJjm+qcZ3tNZOtdY+m+SBbDz+D73GmTcm+YUkmxdeG+V+ZFdyYp/kxFLJiSXVKCdYMhmxADmxNJPLiEROHGSNM0cqJ6bQHLo6yec2XT4x2zYZVXVtkhuSfCjJM1trJ5ONJ/0kV41XWX4jG3+MZ2eXn57ky5seSFO4L5+T5EtJfm82rvrmqro8E7ofW2ufT/Jr2ej4nkzySJKPZHr3ZbLz/TbVx9FPJvnz2deTqbGqXpXk8621j225ajI1cp7J7xc5sRA5sVxyYgnkxEqZ/D6ZcEYkcmJhK5YRiZxYiqOYE1NoDtU22yZzCrWqekqSP07yM621fx67nkFVvTLJQ621j2zevM1Nx74vjyW5McmbWms3JHk00xiffcLsfbY3J7kuyTckuTwb44BbjX1fzjO5fV9Vb8jGSPUfDpu2udmh11hVlyV5Q5L/tt3V22yb8n7vxaT3i5xYmJw4HJPb93KCJZn0PplqRiRyYlmOSEYkE9z3cuJwTaE5dCLJNZsuH0/yhZFqOU9VXZiNJ/M/bK39yWzzF4exsNnnh0Yq73uSvKqq/iEb47MvzUbn/4rZOGMyjfvyRJL+fO1mAAACU0lEQVQTrbUPzS6/KxtP7lO5H5Pk+5N8trX2pdba6SR/kuS7M737Mtn5fpvU46iqbk3yyiQ/1lobngynUuM3ZSO8PzZ7/BxP8ndV9a8ynRo532T3i5xYCjmxXHJicXJitUx2n0w8IxI5sSyrlBGJnFiGI5kTU2gOfTjJ9bWxmvtF2Vhg6s6Raxreb/uWJPe31n5901V3Jrl19vWtSd592LUlSWvt9a214621a7Nxn/11a+3Hkrw/yQ+PXd+gtfaPST5XVc+dbXpZkvsykftx5sEkL66qy2b7fahxUvflzE73251Jfny2Ov6LkzwyjIsetqq6KckvJnlVa+2xTVfdmeSWqrq4qq7LxiJtf3vY9bXWPt5au6q1du3s8XMiyY2zv9XJ3I+cR07sg5xYKjmxRHKCJZMR+yQnlmaVMiKREws7sjnRWhv9I8kPZWMV8v+b5A1j1zOr6XuzMf51T5L/M/v4oWy8D/euJJ+efX7aBGp9SZL3zL5+TjYeIA8k+Z9JLp5Afd+R5O7ZfflnSa6c2v2Y5JeTfDLJJ5L8QZKLx74vk7w9G+9bPp2NJ5zX7HS/ZWN88Xdmj6GPZ+NsCWPV+EA23mc7PG7++6bbv2FW46eSvHysGrdc/w9JnjHm/ehjT/tRTixWq5xYvEY5sbwa5YSPZe9DGbF4vXJisfomlxGzuuTEAdW45fojkRM1+wcAAAAA0KEpvK0MAAAAgJFoDgEAAAB0THMIAAAAoGOaQwAAAAAd0xwCAAAA6JjmEAAAAEDHNIcAAAAAOqY5BAAAANCx/w+0hhXM0JTTaQAAAABJRU5ErkJggg==\n",
      "text/plain": [
       "<Figure size 1440x720 with 3 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "image, label = iter(Pong_train).next()\n",
    "with torch.no_grad():\n",
    "    mu_z1, log_std_z1, mu_x1, log_std_x1, z1, x1 = Pong_VAE1(image[0,:,:].view(-1), noise=False)\n",
    "    mu_z2, log_std_z2, mu_x2, log_std_x2, z2, x2 = Pong_VAE1(image[0,:,:].view(-1), noise=True)\n",
    "x1 = x1.view(160,160)\n",
    "x2 = x2.view(160,160)\n",
    "plt.figure(figsize=(20,10))\n",
    "plt.subplot(1,3,1)\n",
    "plot_image(tensor=image[0,:,:])\n",
    "plt.subplot(1,3,2)\n",
    "plot_image(x1)\n",
    "plt.subplot(1,3,3)\n",
    "plot_image(x2)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train(agent, CNN=False, n_episode, max_t):\n",
    "    score_window = deque(maxlen=100)\n",
    "    all_score = []\n",
    "    for i in range(1, n_episode+1):\n",
    "        if CNN:\n",
    "            previous_state = np.zeros(20)\n",
    "        else\n",
    "            previous_state = np.zeros((1, 160, 160))\n",
    "        current_state = process_image(env.reset())\n",
    "        episodic_score = 0\n",
    "        if CNN:\n",
    "            state = np.concatenate(previous_state, current_state)\n",
    "        else:\n",
    "            previous_state = previous_state.reshape(-1)\n",
    "            with torch.no_grad():\n",
    "                current_state, _, _, _, _, _ = Pong_VAE1(torch.tensor(current_state.reshape(1, -1), dtype=torch.float), noise=False)\n",
    "                current_state = current_state.view(-1).numpy()\n",
    "            state = np.concatenate(previous_state, current_state)\n",
    "        for t in range(max_t):\n",
    "            action = agent.act(state)\n",
    "            next_state, reward, done, _ = env.step(action)\n",
    "            next_state = process_image(next_state)\n",
    "            if CNN:\n",
    "                next_state = np.concatenate(current_state, next_state)\n",
    "            else:\n",
    "                with torch.no_grad():\n",
    "                    next_state, _, _, _, _, _ = Pong_VAE1(torch.tensor(next_state.reshape(1, -1), dtype=torch.float), noise=False)\n",
    "                    next_state = next_state.view(-1).numpy()\n",
    "                next_state = np.concatenate(current_state, next_state)\n",
    "            agent.memory.add(state, action, reward, next_state, done)\n",
    "            episodic_score += reward\n",
    "            previous_state = current_state.copy()\n",
    "            current_state = next_state[1].copy()\n",
    "            if t % 4 ==0:\n",
    "                if len(agent.memory)>agent.memory.batch_size:\n",
    "                    experience = agent.memory.sample()\n",
    "                    agent.learn(experience)\n",
    "            if done:\n",
    "                break\n",
    "        score_window.append(episodic_score)\n",
    "        all_score.append(episodic_score)\n",
    "        print('\\rEpisode {}\\tEpisodic Score: {:.2f}\\tAverage Score: {:.2f}'.format(i_episode, episodic_score, np.mean(scores_window),), end=\"\")\n",
    "        if i_episode % 100 == 0:\n",
    "            print('')\n",
    "        means.append(np.mean(scores_window))\n",
    "            "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for i in range(1, n_episode+1):\n",
    "    state = env.reset()\n",
    "    score = 0\n",
    "    for t in range(max_t):\n",
    "        recoder.append(process_image(state))\n",
    "        action = np.random.randint(0,6)\n",
    "        env.render()\n",
    "        next_state, reward, done, _ = env.step(action)\n",
    "        state = next_state\n",
    "        score += reward\n",
    "        if done:\n",
    "            break \n",
    "    scores.append(score)"
   ]
  }
 ],
 "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.6.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}