gravitar-code.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "# trained on local machine - GTX 1060 3GB GRAM\n",
    "# base code from https://pytorch.org/tutorials/intermediate/reinforcement_q_learning.html#training\n",
    "# RNN inspiration from https://github.com/neka-nat/distributed_rl/tree/master/distributed_rl\n",
    "#                 R2D2 paper: https://openreview.net/pdf?id=r1lyTjAqYX\n",
    "# Extra DQN changes inspired from: https://github.com/fg91/Deep-Q-Learning/blob/master/DQN.ipynb\n",
    "# Hyperparams from rainbow paper - https://arxiv.org/pdf/1710.02298.pdf\n",
    "\n",
    "\n",
    "# https://github.com/openai/baselines/blob/master/baselines/common/atari_wrappers.p\n",
    "# Section containing OpenAI atari wrappers from their repository above\n",
    "\n",
    "import numpy as np\n",
    "import os\n",
    "os.environ.setdefault('PATH', '')\n",
    "from collections import deque\n",
    "import gym\n",
    "from gym import spaces\n",
    "import cv2\n",
    "cv2.ocl.setUseOpenCL(False)\n",
    "\n",
    "\n",
    "\n",
    "class NoopResetEnv(gym.Wrapper):\n",
    "    def __init__(self, env, noop_max=30):\n",
    "        \"\"\"Sample initial states by taking random number of no-ops on reset.\n",
    "        No-op is assumed to be action 0.\n",
    "        \"\"\"\n",
    "        gym.Wrapper.__init__(self, env)\n",
    "        self.noop_max = noop_max\n",
    "        self.override_num_noops = None\n",
    "        self.noop_action = 0\n",
    "        assert env.unwrapped.get_action_meanings()[0] == 'NOOP'\n",
    "\n",
    "    def reset(self, **kwargs):\n",
    "        \"\"\" Do no-op action for a number of steps in [1, noop_max].\"\"\"\n",
    "        self.env.reset(**kwargs)\n",
    "        if self.override_num_noops is not None:\n",
    "            noops = self.override_num_noops\n",
    "        else:\n",
    "            noops = self.unwrapped.np_random.randint(1, self.noop_max + 1) #pylint: disable=E1101\n",
    "        assert noops > 0\n",
    "        obs = None\n",
    "        for _ in range(noops):\n",
    "            obs, _, done, _ = self.env.step(self.noop_action)\n",
    "            if done:\n",
    "                obs = self.env.reset(**kwargs)\n",
    "        return obs\n",
    "\n",
    "    def step(self, ac):\n",
    "        return self.env.step(ac)\n",
    "\n",
    "class FireResetEnv(gym.Wrapper):\n",
    "    def __init__(self, env):\n",
    "        \"\"\"Take action on reset for environments that are fixed until firing.\"\"\"\n",
    "        gym.Wrapper.__init__(self, env)\n",
    "        assert env.unwrapped.get_action_meanings()[1] == 'FIRE'\n",
    "        assert len(env.unwrapped.get_action_meanings()) >= 3\n",
    "\n",
    "    def reset(self, **kwargs):\n",
    "        self.env.reset(**kwargs)\n",
    "        obs, _, done, _ = self.env.step(1)\n",
    "        if done:\n",
    "            self.env.reset(**kwargs)\n",
    "        obs, _, done, _ = self.env.step(2)\n",
    "        if done:\n",
    "            self.env.reset(**kwargs)\n",
    "        return obs\n",
    "\n",
    "    def step(self, ac):\n",
    "        return self.env.step(ac)\n",
    "\n",
    "class EpisodicLifeEnv(gym.Wrapper):\n",
    "    def __init__(self, env):\n",
    "        \"\"\"Make end-of-life == end-of-episode, but only reset on true game over.\n",
    "        Done by DeepMind for the DQN and co. since it helps value estimation.\n",
    "        \"\"\"\n",
    "        gym.Wrapper.__init__(self, env)\n",
    "        self.lives = 0\n",
    "        self.was_real_done  = True\n",
    "\n",
    "    def step(self, action):\n",
    "        obs, reward, done, info = self.env.step(action)\n",
    "        self.was_real_done = done\n",
    "        # check current lives, make loss of life terminal,\n",
    "        # then update lives to handle bonus lives\n",
    "        lives = self.env.unwrapped.ale.lives()\n",
    "        if lives < self.lives and lives > 0:\n",
    "            # for Qbert sometimes we stay in lives == 0 condition for a few frames\n",
    "            # so it's important to keep lives > 0, so that we only reset once\n",
    "            # the environment advertises done.\n",
    "            done = True\n",
    "        self.lives = lives\n",
    "        return obs, reward, done, info\n",
    "\n",
    "    def reset(self, **kwargs):\n",
    "        \"\"\"Reset only when lives are exhausted.\n",
    "        This way all states are still reachable even though lives are episodic,\n",
    "        and the learner need not know about any of this behind-the-scenes.\n",
    "        \"\"\"\n",
    "        if self.was_real_done:\n",
    "            obs = self.env.reset(**kwargs)\n",
    "        else:\n",
    "            # no-op step to advance from terminal/lost life state\n",
    "            obs, _, _, _ = self.env.step(0)\n",
    "        self.lives = self.env.unwrapped.ale.lives()\n",
    "        return obs\n",
    "\n",
    "class MaxAndSkipEnv(gym.Wrapper):\n",
    "    def __init__(self, env, skip=4):\n",
    "        \"\"\"Return only every `skip`-th frame\"\"\"\n",
    "        gym.Wrapper.__init__(self, env)\n",
    "        # most recent raw observations (for max pooling across time steps)\n",
    "        self._obs_buffer = np.zeros((2,)+env.observation_space.shape, dtype=np.uint8)\n",
    "        self._skip       = skip\n",
    "\n",
    "    def step(self, action):\n",
    "        \"\"\"Repeat action, sum reward, and max over last observations.\"\"\"\n",
    "        total_reward = 0.0\n",
    "        done = None\n",
    "        for i in range(self._skip):\n",
    "            obs, reward, done, info = self.env.step(action)\n",
    "            if i == self._skip - 2: self._obs_buffer[0] = obs\n",
    "            if i == self._skip - 1: self._obs_buffer[1] = obs\n",
    "            total_reward += reward\n",
    "            if done:\n",
    "                break\n",
    "        # Note that the observation on the done=True frame\n",
    "        # doesn't matter\n",
    "        max_frame = self._obs_buffer.max(axis=0)\n",
    "\n",
    "        return max_frame, total_reward, done, info\n",
    "\n",
    "    def reset(self, **kwargs):\n",
    "        return self.env.reset(**kwargs)\n",
    "\n",
    "class ClipRewardEnv(gym.RewardWrapper):\n",
    "    def __init__(self, env):\n",
    "        gym.RewardWrapper.__init__(self, env)\n",
    "\n",
    "    def reward(self, reward):\n",
    "        \"\"\"Bin reward to {+1, 0, -1} by its sign.\"\"\"\n",
    "        return np.sign(reward)\n",
    "\n",
    "\n",
    "class WarpFrame(gym.ObservationWrapper):\n",
    "    # def __init__(self, env, width=84, height=84, grayscale=True, dict_space_key=None):\n",
    "    def __init__(self, env, width=84, height=84, grayscale=True, dict_space_key=None):\n",
    "        \"\"\"\n",
    "        Warp frames to 84x84 as done in the Nature paper and later work.\n",
    "        If the environment uses dictionary observations, `dict_space_key` can be specified which indicates which\n",
    "        observation should be warped.\n",
    "        \"\"\"\n",
    "        super().__init__(env)\n",
    "        self._width = width\n",
    "        self._height = height\n",
    "        self._grayscale = grayscale\n",
    "        self._key = dict_space_key\n",
    "        if self._grayscale:\n",
    "            num_colors = 1\n",
    "        else:\n",
    "            num_colors = 3\n",
    "\n",
    "        new_space = gym.spaces.Box(\n",
    "            low=0,\n",
    "            high=255,\n",
    "            shape=(self._height, self._width, num_colors),\n",
    "            dtype=np.uint8,\n",
    "        )\n",
    "        if self._key is None:\n",
    "            original_space = self.observation_space\n",
    "            self.observation_space = new_space\n",
    "        else:\n",
    "            original_space = self.observation_space.spaces[self._key]\n",
    "            self.observation_space.spaces[self._key] = new_space\n",
    "        assert original_space.dtype == np.uint8 and len(original_space.shape) == 3\n",
    "\n",
    "    def observation(self, obs):\n",
    "        if self._key is None:\n",
    "            frame = obs\n",
    "        else:\n",
    "            frame = obs[self._key]\n",
    "\n",
    "        if self._grayscale:\n",
    "            frame = cv2.cvtColor(frame, cv2.COLOR_RGB2GRAY)\n",
    "        frame = cv2.resize(\n",
    "            frame, (self._width, self._height), interpolation=cv2.INTER_AREA\n",
    "        )\n",
    "        if self._grayscale:\n",
    "            frame = np.expand_dims(frame, -1)\n",
    "\n",
    "        if self._key is None:\n",
    "            obs = frame\n",
    "        else:\n",
    "            obs = obs.copy()\n",
    "            obs[self._key] = frame\n",
    "        return obs\n",
    "\n",
    "\n",
    "class FrameStack(gym.Wrapper):\n",
    "    def __init__(self, env, k):\n",
    "        \"\"\"Stack k last frames.\n",
    "        Returns lazy array, which is much more memory efficient.\n",
    "        See Also\n",
    "        --------\n",
    "        baselines.common.atari_wrappers.LazyFrames\n",
    "        \"\"\"\n",
    "        gym.Wrapper.__init__(self, env)\n",
    "        self.k = k\n",
    "        self.frames = deque([], maxlen=k)\n",
    "        shp = env.observation_space.shape\n",
    "        self.observation_space = spaces.Box(low=0, high=255, shape=(shp[:-1] + (shp[-1] * k,)), dtype=env.observation_space.dtype)\n",
    "\n",
    "    def reset(self):\n",
    "        ob = self.env.reset()\n",
    "        for _ in range(self.k):\n",
    "            self.frames.append(ob)\n",
    "        return self._get_ob()\n",
    "\n",
    "    def step(self, action):\n",
    "        ob, reward, done, info = self.env.step(action)\n",
    "        self.frames.append(ob)\n",
    "        return self._get_ob(), reward, done, info\n",
    "\n",
    "    def _get_ob(self):\n",
    "        assert len(self.frames) == self.k\n",
    "        return LazyFrames(list(self.frames))\n",
    "\n",
    "class ScaledFloatFrame(gym.ObservationWrapper):\n",
    "    def __init__(self, env):\n",
    "        gym.ObservationWrapper.__init__(self, env)\n",
    "        self.observation_space = gym.spaces.Box(low=0, high=1, shape=env.observation_space.shape, dtype=np.float32)\n",
    "\n",
    "    def observation(self, observation):\n",
    "        # careful! This undoes the memory optimization, use\n",
    "        # with smaller replay buffers only.\n",
    "        return np.array(observation).astype(np.float32) / 255.0\n",
    "\n",
    "class LazyFrames(object):\n",
    "    def __init__(self, frames):\n",
    "        \"\"\"This object ensures that common frames between the observations are only stored once.\n",
    "        It exists purely to optimize memory usage which can be huge for DQN's 1M frames replay\n",
    "        buffers.\n",
    "        This object should only be converted to numpy array before being passed to the model.\n",
    "        You'd not believe how complex the previous solution was.\"\"\"\n",
    "        self._frames = frames\n",
    "        self._out = None\n",
    "\n",
    "    def _force(self):\n",
    "        if self._out is None:\n",
    "            self._out = np.concatenate(self._frames, axis=-1)\n",
    "            self._frames = None\n",
    "        return self._out\n",
    "\n",
    "    def __array__(self, dtype=None):\n",
    "        out = self._force()\n",
    "        if dtype is not None:\n",
    "            out = out.astype(dtype)\n",
    "        return out\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self._force())\n",
    "\n",
    "    def __getitem__(self, i):\n",
    "        return self._force()[i]\n",
    "\n",
    "    def count(self):\n",
    "        frames = self._force()\n",
    "        return frames.shape[frames.ndim - 1]\n",
    "\n",
    "    def frame(self, i):\n",
    "        return self._force()[..., i]\n",
    "\n",
    "def make_atari(env_id,video_every, max_episode_steps=None):\n",
    "    env = gym.make(env_id)\n",
    "\n",
    "    # assert 'NoFrameskip' in env.spec.id\n",
    "    env = NoopResetEnv(env, noop_max=30)\n",
    "    env = MaxAndSkipEnv(env, skip=4)\n",
    "    if max_episode_steps is not None:\n",
    "        pass\n",
    "        # env = TimeLimit(env, max_episode_steps=max_episode_steps)\n",
    "    # Added video Monitor\n",
    "    if video_every!=0:\n",
    "        env = gym.wrappers.Monitor(env, \"./recordings/\" + env_id + '/', video_callable=lambda episode_id: (episode_id % video_every) == 0, force=True)\n",
    "    return env\n",
    "\n",
    "def wrap_deepmind(env, episode_life=True, clip_rewards=True, frame_stack=False, scale=False):\n",
    "    \"\"\"Configure environment for DeepMind-style Atari.\n",
    "    \"\"\"\n",
    "    if episode_life:\n",
    "        env = EpisodicLifeEnv(env)\n",
    "    if 'FIRE' in env.unwrapped.get_action_meanings():\n",
    "        env = FireResetEnv(env)\n",
    "    env = WarpFrame(env)\n",
    "    if scale:\n",
    "        env = ScaledFloatFrame(env)\n",
    "    if clip_rewards:\n",
    "        env = ClipRewardEnv(env)\n",
    "    if frame_stack:\n",
    "        env = FrameStack(env, 4)\n",
    "    return env\n"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "\n",
    "import gym\n",
    "import gym.wrappers.monitoring.video_recorder\n",
    "\n",
    "import math\n",
    "import random\n",
    "import matplotlib\n",
    "import matplotlib.pyplot as plt\n",
    "from collections import namedtuple\n",
    "from  collections import deque\n",
    "from collections import Counter\n",
    "from PIL import Image\n",
    "\n",
    "# import atari_wrappers\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "import torch.nn.functional as F\n",
    "import torchvision.transforms as T\n",
    "from tqdm.auto import tqdm\n",
    "\n",
    "torch.cuda.empty_cache()"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "# env = gym.make('Breakout-v4')#.unwrapped\n",
    "# env = gym.make('Gravitar-v0').unwrapped\n",
    "# game = 'BreakoutNoFrameskip-v4'\n",
    "game = 'GravitarNoFrameskip-v4'\n",
    "# game= 'PongNoFrameskip-v0'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "# Wrappers\n",
    "video_every=5\n",
    "env = make_atari(game,video_every,max_episode_steps=None)\n",
    "# env = gym.wrappers.Monitor(env, \"./recordings/\"+game+'/', video_callable=lambda episode_id: (episode_id%video_every)==0,force=True)\n",
    "# env = atari_wrappers.wrap_deepmind(env,episode_life=True,clip_rewards=False,frame_stack=True)\n",
    "env = wrap_deepmind(env,episode_life=False,clip_rewards=False,frame_stack=True)\n",
    "\n",
    "# set up matplotlib\n",
    "is_ipython = 'inline' in matplotlib.get_backend()\n",
    "if is_ipython:\n",
    "    from IPython import display\n",
    "\n",
    "plt.ion()\n",
    "\n",
    "# if gpu is to be used\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "# Replay memory class, storing states, actions, following state and the reward gained\n",
    "\n",
    "Transition = namedtuple('Transition',\n",
    "                        ('state', 'action', 'next_state', 'reward', 'done'))\n",
    "Sequence = namedtuple('Sequence',('transitions','recurrent_state'))\n",
    "\n",
    "class ReplayMemory(object):\n",
    "\n",
    "    def __init__(self, capacity, files):\n",
    "        self.files = files\n",
    "        self.currNewestFile=0\n",
    "        self.fileName='./memory/memory-data'#+file+'.pt'\n",
    "        self.liveFiles = set()\n",
    "        self.capacity = capacity\n",
    "        self.memory = []\n",
    "        self.position = 0\n",
    "        self.priorities = np.zeros(capacity)\n",
    "        self.alpha = 0.6\n",
    "        load = False\n",
    "        if load:\n",
    "            self.load()\n",
    "\n",
    "\n",
    "    def push(self, seq):\n",
    "        \"\"\"Saves a transition.\"\"\"\n",
    "        if len(self.memory) < self.capacity:\n",
    "            self.memory.append(None)\n",
    "        self.memory[self.position] = seq\n",
    "        self.priorities[self.position] = 1\n",
    "        self.position = (self.position + 1) % self.capacity\n",
    "        if self.position==0:\n",
    "            self.save()\n",
    "\n",
    "    def save(self):\n",
    "\n",
    "        try:\n",
    "            for i in range(0,capacity,1000):\n",
    "                memToSave = {'mem': self.memory[i:i+1000],'pri':self.priorities[i:i+1000]}\n",
    "                torch.save(memToSave,self.fileName+str(i)+'.pt')\n",
    "        except:\n",
    "            print('Error While Saving')\n",
    "            pass\n",
    "    def load(self):\n",
    "        try:\n",
    "            for i in range(0,capacity,1000):\n",
    "                TensorDict = torch.load(self.fileName+str(i)+'.pt')\n",
    "                self.memory+=TensorDict['mem']\n",
    "                self.priorities[i:i+1000] = TensorDict['pri']\n",
    "        except:\n",
    "            print('Error While Saving')\n",
    "            pass\n",
    "\n",
    "        self.position = len(self.memory) % self.capacity\n",
    "\n",
    "    def memSample(self):\n",
    "        probs = np.array(list(range(1,len(self.liveFiles)+2)))\n",
    "        probs[-1]*=2\n",
    "        probs= probs / probs.sum()\n",
    "\n",
    "        experiences = np.random.choice(len(self.liveFiles)+1,BATCH_SIZE,p=probs)\n",
    "        experiences = self.calcFreq(experiences)\n",
    "\n",
    "        data = []\n",
    "        for i in experiences:\n",
    "            # experiences[i] (batch size from that file), i (index of file)\n",
    "            if i==len(self.liveFiles):\n",
    "                data+= self.sample(experiences[i])\n",
    "            else:\n",
    "                x = (self.currNewestFile-i-1)%self.files\n",
    "                data+= self.loadFile(self.fileName+str(x)+'.pt',experiences[i])\n",
    "        return data\n",
    "\n",
    "\n",
    "    def calcFreq(self,data):\n",
    "        freq = {}\n",
    "        for i in data:\n",
    "            if i in freq:\n",
    "                freq[i]+=1\n",
    "            else:\n",
    "                freq[i]=1\n",
    "        return freq\n",
    "\n",
    "    def loadFile(self,name,batch_size):\n",
    "        return random.sample(torch.load(name),batch_size)\n",
    "\n",
    "    def sample(self, batch_size):\n",
    "        return random.sample(self.memory, batch_size)\n",
    "\n",
    "    def PERSample(self,batch_size, beta =0.6):\n",
    "        # beta changes over time\n",
    "        #Can change to use sum-tree\n",
    "        if len(self.memory) < self.capacity:\n",
    "            tmpPriorities = self.priorities[:self.position]\n",
    "        else:\n",
    "            tmpPriorities = self.priorities\n",
    "\n",
    "        #priority_j ^alpha / sum priorities\n",
    "        probs = tmpPriorities**self.alpha\n",
    "        probs /= probs.sum()\n",
    "\n",
    "        indices = np.random.choice(len(self.memory), batch_size, p=probs)\n",
    "        tmpPriorities = tmpPriorities[indices]\n",
    "        loss_weights = (len(tmpPriorities)*tmpPriorities/len(self.memory))**(-beta)\n",
    "        loss_weights /= loss_weights.max()\n",
    "\n",
    "        return [self.memory[index] for index in indices], indices, loss_weights\n",
    "\n",
    "    def updatePriorities(self, indices, newPriorities):\n",
    "        for i in range(len(indices)):\n",
    "            self.priorities[indices[i]] = newPriorities[i]\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.memory)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "# https://github.com/Shmuma/ptan/blob/master/samples/rainbow/lib/dqn_model.py\n",
    "# Factorised Noisy Layer\n",
    "class NoisyFactorizedLinear(nn.Linear):\n",
    "    \"\"\"\n",
    "    NoisyNet layer with factorized gaussian noise\n",
    "    N.B. nn.Linear already initializes weight and bias to\n",
    "    \"\"\"\n",
    "    def __init__(self, in_features, out_features, sigma_zero=0.4, bias=True):\n",
    "        super(NoisyFactorizedLinear, self).__init__(in_features, out_features, bias=bias)\n",
    "        sigma_init = sigma_zero / math.sqrt(in_features)\n",
    "        self.sigma_weight = nn.Parameter(torch.Tensor(out_features, in_features).fill_(sigma_init))\n",
    "        self.register_buffer(\"epsilon_input\", torch.zeros(1, in_features))\n",
    "        self.register_buffer(\"epsilon_output\", torch.zeros(out_features, 1))\n",
    "        if bias:\n",
    "            self.sigma_bias = nn.Parameter(torch.Tensor(out_features).fill_(sigma_init))\n",
    "\n",
    "    def forward(self, input):\n",
    "        torch.randn(self.epsilon_input.size(), out=self.epsilon_input)\n",
    "        torch.randn(self.epsilon_output.size(), out=self.epsilon_output)\n",
    "\n",
    "        func = lambda x: torch.sign(x) * torch.sqrt(torch.abs(x))\n",
    "        eps_in = func(self.epsilon_input)\n",
    "        eps_out = func(self.epsilon_output)\n",
    "\n",
    "        bias = self.bias\n",
    "        if bias is not None:\n",
    "            bias = bias + self.sigma_bias * torch.autograd.Variable(eps_out.t())\n",
    "        noise_v = torch.autograd.Variable(torch.mul(eps_in, eps_out))\n",
    "        return F.linear(input, self.weight + self.sigma_weight * noise_v, bias)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "# Number of Linear input connections depends on output of conv2d layers\n",
    "# and therefore the input image size, so compute it.\n",
    "def conv2d_size_out(size, kernel_size = 5, stride = 2):\n",
    "    return (size - (kernel_size - 1) - 1) // stride  + 1\n",
    "\n",
    "\n",
    "class DuelingLSTMDQN(nn.Module):\n",
    "\n",
    "    def __init__(self, h, w, outputs):\n",
    "        super(DuelingLSTMDQN, self).__init__()\n",
    "        if scaledScreenSize[1] > 80:\n",
    "            conv_out = 64\n",
    "            # Hyper parameters from Rainbow paper - https://arxiv.org/pdf/1710.02298.pdf\n",
    "            if grey:\n",
    "                self.conv1 = nn.Conv2d(1*4, 32, kernel_size=8, stride=4)\n",
    "            else:\n",
    "                self.conv1 = nn.Conv2d(3*4, 32, kernel_size=8, stride=4)\n",
    "            self.bn1 = nn.BatchNorm2d(32)\n",
    "            self.conv2 = nn.Conv2d(32, 64, kernel_size=4, stride=2)\n",
    "            self.bn2 = nn.BatchNorm2d(64)\n",
    "            self.conv3 = nn.Conv2d(64, conv_out, kernel_size=3, stride=1)\n",
    "            self.bn3 = nn.BatchNorm2d(conv_out)\n",
    "            convw = conv2d_size_out(conv2d_size_out(conv2d_size_out(w,kernel_size=8,stride=4),kernel_size=4),kernel_size=3,stride=1)\n",
    "            convh = conv2d_size_out(conv2d_size_out(conv2d_size_out(h,kernel_size=8,stride=4),kernel_size=4),kernel_size=3,stride=1)\n",
    "\n",
    "        else: # scaledScreenSize[0] == 64:\n",
    "            # for reduced model complexity\n",
    "            conv_out = 32\n",
    "            if grey:\n",
    "                self.conv1 = nn.Conv2d(1*4, 16, kernel_size=6, stride=3)\n",
    "            else:\n",
    "                self.conv1 = nn.Conv2d(3*4, 16, kernel_size=6, stride=3)\n",
    "            self.bn1 = nn.BatchNorm2d(16)\n",
    "            self.conv2 = nn.Conv2d(16, 32, kernel_size=4, stride=2)\n",
    "            self.bn2 = nn.BatchNorm2d(32)\n",
    "            self.conv3 = nn.Conv2d(32, conv_out, kernel_size=3, stride=1)\n",
    "            self.bn3 = nn.BatchNorm2d(conv_out)\n",
    "            convw = conv2d_size_out(conv2d_size_out(conv2d_size_out(w,kernel_size=6,stride=3),kernel_size=4),kernel_size=3,stride=1)\n",
    "            convh = conv2d_size_out(conv2d_size_out(conv2d_size_out(h,kernel_size=6,stride=3),kernel_size=4),kernel_size=3,stride=1)\n",
    "\n",
    "\n",
    "        lstm_input_size = convw * convh * conv_out\n",
    "        self.linear_input_size = 512\n",
    "        # linear_input_size = convw * convh * conv_out\n",
    "\n",
    "\n",
    "        self.lstm = nn.LSTMCell(lstm_input_size, self.linear_input_size)\n",
    "        self.hiddenStatesBatch = torch.zeros(BATCH_SIZE,self.linear_input_size).to(device)\n",
    "        self.cellStatesBatch = torch.zeros(BATCH_SIZE,self.linear_input_size).to(device)\n",
    "        self.hiddenStatesLive = torch.zeros(1,self.linear_input_size).to(device)\n",
    "        self.cellStatesLive = torch.zeros(1,self.linear_input_size).to(device)\n",
    "\n",
    "        # dualing networks\n",
    "        # https://arxiv.org/pdf/1511.06581.pdf\n",
    "        # self.adv1 = NoisyFactorizedLinear(self.linear_input_size, 512)\n",
    "        # self.adv2 = NoisyFactorizedLinear(512, outputs)\n",
    "        # self.val1 = NoisyFactorizedLinear(self.linear_input_size, 512)\n",
    "        # self.val2 = NoisyFactorizedLinear(512, 1)\n",
    "\n",
    "        self.adv1 = nn.Linear(self.linear_input_size, 512)\n",
    "        self.adv2 = nn.Linear(512, outputs)\n",
    "        self.val1 = nn.Linear(self.linear_input_size, 512)\n",
    "        self.val2 = nn.Linear(512, 1)\n",
    "        # weight initialisation https://towardsdatascience.com/tutorial-double-deep-q-learning-with-dueling-network-architectures-4c1b3fb7f756\n",
    "        nn.init.kaiming_normal_(self.adv1.weight)\n",
    "        nn.init.kaiming_normal_(self.val1.weight)\n",
    "\n",
    "    def reset(self,done, batch=False):\n",
    "        if batch:\n",
    "            self.hiddenStatesBatch.detach()\n",
    "            self.cellStatesBatch.detach()\n",
    "            if done:\n",
    "                self.hiddenStatesBatch = self.hiddenStatesBatch.zero_()\n",
    "                self.cellStatesBatch = self.hiddenStatesBatch.zero_()\n",
    "        else:\n",
    "            self.hiddenStatesLive.detach()\n",
    "            self.cellStatesLive.detach()\n",
    "            if done:\n",
    "                self.hiddenStatesLive = self.hiddenStatesLive.zero_()\n",
    "                self.cellStatesLive = self.hiddenStatesLive.zero_()\n",
    "\n",
    "    def getState(self, batch=False):\n",
    "        if batch:\n",
    "            return self.hiddenStatesBatch.detach().clone().cpu(), self.cellStatesBatch.detach().clone().cpu()\n",
    "        else:\n",
    "            return self.hiddenStatesLive.detach().clone().cpu(), self.cellStatesLive.detach().clone().cpu()\n",
    "\n",
    "    def setState(self, state,device, batch=False):\n",
    "        if batch:\n",
    "            hiddens, cells = state\n",
    "            self.hiddenStatesBatch = hiddens.clone().to(device)\n",
    "            self.cellStatesBatch = cells.clone().to(device)\n",
    "        else:\n",
    "            hiddens, cells = state\n",
    "            self.hiddenStatesLive = hiddens.clone().to(device)\n",
    "            self.cellStatesLive = cells.clone().to(device)\n",
    "\n",
    "    # Called with either one element to determine next action, or a batch\n",
    "    # during optimisation.\n",
    "    def forward(self, x, hiddenState=None, cellState=None):\n",
    "        x = F.relu(self.bn1(self.conv1(x)))\n",
    "        x = F.relu(self.bn2(self.conv2(x)))\n",
    "        x = F.relu(self.bn3(self.conv3(x)))\n",
    "        x = x.contiguous().view(x.size(0), -1)\n",
    "        if not(hiddenState==None):\n",
    "\n",
    "            self.hiddenStatesLive, self.cellStatesLive = self.lstm(x,(hiddenState,cellState))\n",
    "            adv = F.relu(self.adv1(self.hiddenStatesLive))\n",
    "            val = F.relu(self.val1(self.hiddenStatesLive))\n",
    "\n",
    "        else:\n",
    "            self.hiddenStatesBatch, self.cellStatesBatch = self.lstm(x,(self.hiddenStatesBatch,self.cellStatesBatch))\n",
    "            adv = F.relu(self.adv1(self.hiddenStatesBatch))\n",
    "            val = F.relu(self.val1(self.hiddenStatesBatch))\n",
    "        # Duelling architecture\n",
    "        adv = self.adv2(adv)\n",
    "        val = self.val2(val)\n",
    "        return val + adv - adv.mean(1, keepdim=True)\n",
    "\n",
    "    def calcLoss(self,batch):\n",
    "        copy = DuelingLSTMDQN(screen_height, screen_width, n_actions).to(device)\n",
    "        copy.load_state_dict(self.state_dict())\n",
    "        copy.eval()\n",
    "        batch = Sequence(*zip(*batch))\n",
    "        batch = Sequence(list(zip(*batch.transitions)),list(zip(*batch.recurrent_state)))\n",
    "\n",
    "        hiddenStates = torch.cat(batch.recurrent_state[0])\n",
    "        cellStates = torch.cat(batch.recurrent_state[1])\n",
    "\n",
    "\n",
    "        self.setState((hiddenStates,cellStates),device,True)\n",
    "        target_net.setState((hiddenStates,cellStates),device,True)\n",
    "\n",
    "        #burn in the model\n",
    "        with torch.no_grad():\n",
    "            for t in range(n_burn_in):\n",
    "                tmpTransition = Transition(*zip(*batch.transitions[t]))\n",
    "                state = get_screen(tmpTransition.state,True).to(device)\n",
    "                self.forward(state)\n",
    "                target_net.forward(state)\n",
    "            copy.setState(self.getState(True),device,True)\n",
    "            for t in range(n_burn_in,n_burn_in+n_step_return):\n",
    "                tmpTransition = Transition(*zip(*batch.transitions[t]))\n",
    "                state = get_screen(tmpTransition.state,True).to(device)\n",
    "                copy.forward(state)\n",
    "                target_net.forward(state)\n",
    "        self.reset(False,True)\n",
    "\n",
    "        # init loss array\n",
    "        delta = torch.zeros(len(batch.transitions)-n_burn_in-n_step_return,BATCH_SIZE,1,device=device)\n",
    "        with torch.set_grad_enabled(True):\n",
    "            for t in range(n_burn_in,len(batch.transitions)-n_step_return):\n",
    "                tmpTransition = Transition(*zip(*batch.transitions[t]))\n",
    "                nextTransition = Transition(*zip(*batch.transitions[t+n_step_return]))\n",
    "                state_batch = get_screen(tmpTransition.state,True).to(device)\n",
    "                action_batch = torch.cat(tmpTransition.action).to(device)\n",
    "                reward_batch = torch.stack(tmpTransition.reward).to(device)\n",
    "                next_state_batch = get_screen(nextTransition.state,True).to(device)\n",
    "                done_batch = torch.stack(nextTransition.done).to(device)\n",
    "\n",
    "                state_action_values = self.forward(state_batch)\n",
    "                state_action_values = state_action_values.gather(1,action_batch)\n",
    "                actionPrediction = copy.forward(next_state_batch).argmax(dim=1).unsqueeze(1).detach()\n",
    "\n",
    "                next_state_values = target_net(next_state_batch).gather(1, actionPrediction).detach()\n",
    "\n",
    "                expected_state_action_values = (self.inv_value_function(next_state_values) * GAMMA * (1.0 - done_batch)) \\\n",
    "                                           + reward_batch\n",
    "                expected_state_action_values = self.value_function(expected_state_action_values)\n",
    "\n",
    "                # delta[t-n_burn_in] = F.smooth_l1_loss(state_action_values,expected_state_action_values.float(),reduction='none')\n",
    "                delta[t-n_burn_in] = F.mse_loss(state_action_values,expected_state_action_values.float(),reduction='none')\n",
    "\n",
    "        # 0.9 is eta value in R2 paper\n",
    "        priorities = 0.9 * delta.max(dim=0)[0] + (1.0 - 0.9) * delta.mean(dim=0)\n",
    "        return delta.sum(dim=0), priorities.detach()\n",
    "        # return delta.pow(2).sum(dim=0), priorities.detach()\n",
    "\n",
    "    # https://arxiv.org/pdf/1805.11593.pdf\n",
    "    def value_function(self,x,eps=1e-3):\n",
    "        return x.sign() * ((x.abs()+1).sqrt()-1)+eps*x\n",
    "    def inv_value_function(self,x,eps=1e-3):\n",
    "        return x.sign() * ((((1+4*eps*(x.abs()+1+eps)).sqrt()-1) / (2 *eps)).pow(2)-1)\n",
    "\n"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "scaledScreenSize = (84,84)\n",
    "grey=True\n",
    "\n",
    "def get_screen(screen,batch=False):\n",
    "\n",
    "    temp = np.ascontiguousarray(screen, dtype=np.float32) / 255\n",
    "    # BATCH,x,y,frame\n",
    "    if batch:\n",
    "        return torch.from_numpy(temp).permute(0,3,1,2)\n",
    "    else:\n",
    "        return torch.from_numpy(temp).unsqueeze(0).permute(0,3,1,2)\n",
    "\n",
    "# output example image\n",
    "\n",
    "# temp = env.reset()\n",
    "# env.step(0)\n",
    "# env.step(0)\n",
    "# plt.figure()\n",
    "# if grey:\n",
    "#     plt.imshow(get_screen(temp).cpu().permute(2,0,1)[0].numpy(),\n",
    "#             interpolation='none', cmap='gray')\n",
    "# else:\n",
    "#     plt.imshow(get_screen(temp).cpu().permute(1, 2, 0).numpy(),\n",
    "#             interpolation='none')\n",
    "# plt.title('What the model Sees')\n",
    "# plt.show()\n",
    "#\n",
    "# exit()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "num_episodes = 10000\n",
    "BATCH_SIZE = 23 # small batch due to training on GTX 1060 3GB\n",
    "GAMMA = 0.999**3\n",
    "TARGET_UPDATE = 1000 # frames\n",
    "POLICY_UPDATE = 24 # frames\n",
    "BETA_DECAY = 2000\n",
    "OUTPUT_FREQUENCY = 10\n",
    "n_step_return = 5\n",
    "n_sequence = 50 #size of LSTM\n",
    "n_overlap = 25 #the next RNN state is half of the sequence length\n",
    "n_burn_in = 25 #how many steps before RNN is set to actually update - ie gets better initial state\n",
    "\n",
    "# update to 20 batch - 50 n_sequence, eps decay 150, eps max 0.4, episodic end of life = true -\n",
    "# already made lr and loss power changes\n",
    "\n",
    "# init with correct dims\n",
    "if scaledScreenSize[1]==84:\n",
    "    screen_width = 84\n",
    "    screen_height = 84\n",
    "else:\n",
    "    screen_width = 64\n",
    "    screen_height = 64\n",
    "\n",
    "# Get number of actions from gym action space\n",
    "n_actions = env.action_space.n\n",
    "\n",
    "# Actor generates experiences\n",
    "# Learner learns from experiences every episode\n",
    "# params = torch.load('./memory/target_net.chkpt')\n",
    "Actor_net = DuelingLSTMDQN(screen_height, screen_width, n_actions).to(device)\n",
    "Learner_net = DuelingLSTMDQN(screen_height, screen_width, n_actions).to(device)\n",
    "# Learner_net.load_state_dict(params['A'])\n",
    "target_net = DuelingLSTMDQN(screen_height, screen_width, n_actions).to(device)\n",
    "# target_net.load_state_dict(Learner_net.state_dict())\n",
    "target_net.eval()\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "optimiser = optim.Adam(Learner_net.parameters(),lr=1e-4,eps=1e-3)\n",
    "# optimiser.load_state_dict(params['optimiser'])\n",
    "capacity = 10000\n",
    "files = 5\n",
    "memory = ReplayMemory(capacity=capacity,files=files)\n",
    "\n",
    "\n",
    "def select_action(state,n_episode, LSTMState):\n",
    "\n",
    "\n",
    "    # global steps_done\n",
    "    sample = random.random()\n",
    "    # custom eps greedy function\n",
    "    epsilon = max( abs(math.exp(-(n_episode/400))*0.4*math.sin(n_episode/10)), 0.001)\n",
    "\n",
    "    if sample > epsilon:\n",
    "        with torch.no_grad():\n",
    "            x = Actor_net(get_screen(state).to(device),LSTMState[0].to(device),LSTMState[1].to(device))\n",
    "            return x.max(1)[1].view(1, 1).cpu()\n",
    "    else:\n",
    "        return torch.tensor([[random.randrange(n_actions)]], device='cpu', dtype=torch.long)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "\n",
    "def optimize_model(n_episode):\n",
    "\n",
    "    # beta = 0.75\n",
    "    beta = min(1, 0.4 + 1*(n_episode/BETA_DECAY))\n",
    "\n",
    "    if len(memory) < BATCH_SIZE*2:\n",
    "        return\n",
    "    lossArr = []\n",
    "    for _ in range(2):\n",
    "\n",
    "        Learner_net.reset(done=True,batch=True)\n",
    "\n",
    "\n",
    "        # transitions = memory.sample(BATCH_SIZE)\n",
    "        transitions, indices, loss_weights = memory.PERSample(BATCH_SIZE,beta)\n",
    "        # transitions, indices = memory.PERSample(BATCH_SIZE)\n",
    "        # transitions = memory.memSample()\n",
    "\n",
    "        delta,priorities = Learner_net.calcLoss(transitions)\n",
    "\n",
    "        memory.updatePriorities(indices,priorities)\n",
    "        loss = (delta*torch.Tensor(loss_weights).to(device)).mean().float()\n",
    "        # loss = (delta).mean().float()\n",
    "        lossArr.append(loss.clone().detach().cpu().numpy())\n",
    "        # Optimise the model\n",
    "        optimiser.zero_grad()\n",
    "        loss.backward()\n",
    "        # clip grads\n",
    "        for param in Learner_net.parameters():\n",
    "            param.grad.data.clamp_(-1, 1)\n",
    "        optimiser.step()\n",
    "    return lossArr"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "import time\n",
    "#Training\n",
    "start=time.time()\n",
    "seed = 742\n",
    "torch.manual_seed(seed)\n",
    "env.seed(seed)\n",
    "random.seed(seed)\n",
    "np.random.seed(seed)\n",
    "env.action_space.seed(seed)\n",
    "\n",
    "marking=[]\n",
    "episodeScores=[]\n",
    "losses = []\n",
    "stepBuffer = deque(maxlen=n_step_return)\n",
    "#precalculate GAMMA n step return values\n",
    "gamma_nsteps = [GAMMA ** i for i in range(n_step_return+1)]\n",
    "sequenceBuffer = []\n",
    "recurrentBuffer = []\n",
    "frames = 0\n",
    "n_episode=2300\n",
    "\n",
    "score=0\n",
    "actorUpdates=0\n",
    "# Initialize the environment and state\n",
    "maxLives = env.ale.lives()\n",
    "state = env.reset()\n",
    "lives = maxLives\n",
    "with tqdm() as pbar:\n",
    "    while n_episode<num_episodes:\n",
    "    # for n_episode in tqdm(range(num_episodes+1)):\n",
    "\n",
    "        sequenceBuffer = []\n",
    "        recurrentBuffer = []\n",
    "        while True:\n",
    "\n",
    "            recurrentBuffer.append(Actor_net.getState())\n",
    "            action = select_action(state,n_episode,recurrentBuffer[-1])\n",
    "\n",
    "            obvs, reward, done, info = env.step(action.item())\n",
    "            score+=reward\n",
    "\n",
    "            # Reward clipping and negative reward on death\n",
    "            if info['ale.lives'] != lives:\n",
    "                pbar.update(1)\n",
    "                pbar.set_postfix({'episode':n_episode,'score':score,'Actor Updates':actorUpdates,'Frames':frames})\n",
    "                reward-=1\n",
    "                lives = info['ale.lives']\n",
    "            # reward = np.clip(reward, -1, 1)\n",
    "\n",
    "            reward = torch.tensor([reward])\n",
    "\n",
    "            # Observe new state\n",
    "            next_state = obvs\n",
    "            if done:\n",
    "                next_state = None\n",
    "\n",
    "            stepBuffer.append(Transition(state,action,next_state,reward,torch.tensor([float(done)])))\n",
    "            if len(stepBuffer) == stepBuffer.maxlen:\n",
    "                reward_nstep = sum([gamma_nsteps[-(i+2)] * stepBuffer[i].reward for i in range(n_step_return)])\n",
    "                sequenceBuffer.append(Transition(stepBuffer[0].state,stepBuffer[0].action,stepBuffer[0].next_state,reward_nstep,torch.tensor([float(done)])))\n",
    "            if (len (sequenceBuffer)) == n_sequence+n_burn_in:\n",
    "                memory.push(Sequence(sequenceBuffer,recurrentBuffer[0]))\n",
    "                #chop sequence and recurrent buffers\n",
    "                sequenceBuffer = sequenceBuffer[-(n_burn_in+n_overlap):]\n",
    "                recurrentBuffer = recurrentBuffer[-(n_burn_in+n_overlap+n_step_return):]\n",
    "            elif done and len(sequenceBuffer) > n_burn_in+n_overlap+n_step_return:\n",
    "                # case of we're done so we push the rest of the sequence to memory\n",
    "                memory.push(Sequence(sequenceBuffer,recurrentBuffer[0]))\n",
    "            if frames % POLICY_UPDATE == 0:\n",
    "                x = optimize_model(n_episode)\n",
    "                if x != None:\n",
    "                    losses += x\n",
    "            # Update the target network, copying all weights and biases in DQN\n",
    "            if frames % TARGET_UPDATE == 0:\n",
    "                actorUpdates+=1\n",
    "                Actor_net.load_state_dict(Learner_net.state_dict())\n",
    "                target_net.load_state_dict(Learner_net.state_dict())\n",
    "\n",
    "            # Move to the next state\n",
    "            state = next_state\n",
    "            frames+=1\n",
    "            if done:\n",
    "                stepBuffer.clear()\n",
    "\n",
    "                break\n",
    "\n",
    "        state = env.reset()\n",
    "        lives = env.ale.lives()\n",
    "        if lives == maxLives:\n",
    "            episodeScores.append(int(score))\n",
    "            marking.append(score)\n",
    "\n",
    "            # do not change lines 44-48 here, they are for marking the submission log\n",
    "            if n_episode%100 == 0:\n",
    "                print(\"marking, episode: {}, score: {:.1f}, mean_score: {:.2f}, std_score: {:.2f}\".format(\n",
    "                    n_episode, score, np.array(marking).mean(), np.array(marking).std()))\n",
    "                marking = []\n",
    "\n",
    "            if n_episode % OUTPUT_FREQUENCY == 0:\n",
    "                if len(marking)>0:\n",
    "                    print(\"episode: {}, mean_last30: {:.2f}, mean_score: {:.2f}, std_score: {:.2f}, mean_loss: {:.4f}\".format(\n",
    "                        n_episode, np.array(episodeScores[-30:]).mean(), np.array(marking).mean(), np.array(marking).std(),np.array(losses[-20:]).mean()),episodeScores[-OUTPUT_FREQUENCY:])\n",
    "            n_episode+=1\n",
    "            score=0\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "print('Complete')\n",
    "# env.render()\n",
    "env.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "memory.save()\n",
    "torch.save({'A':target_net.state_dict(), 'optimiser':optimiser.state_dict(),}, './memory/target_net.chkpt')"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "# Trend Graphs\n",
    "episodeScoresEnum = sorted(enumerate(episodeScores),reverse=True,key=lambda x:x[1])\n",
    "print('Top 5 scores overall')\n",
    "print(episodeScoresEnum[:5])\n",
    "tempScores = sorted([(i,j) for i,j in enumerate(episodeScores) if i%5==0],key=lambda x:x[1],reverse=True)\n",
    "print('Top videos:')\n",
    "print(tempScores[:10])\n",
    "print('Bottom Videos')\n",
    "print(tempScores[-10:][::-1])\n",
    "print('Total Frames Trained on: ',frames)\n",
    "\n"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "averages30 = []\n",
    "averages = []\n",
    "\n",
    "for i in range(len(losses)):\n",
    "    averages.append(np.array(losses[max(i-100,0):i+1]).mean())\n",
    "    averages30.append(np.array(losses[max(i-30,0):i+1]).mean())\n",
    "\n",
    "plt.plot(range(len(losses)),losses,alpha=0.4,label='actual')\n",
    "plt.plot(range(len(averages)),averages,label='100 average')\n",
    "plt.plot(range(len(averages30)),averages30,label='30 average',alpha=0.6)\n",
    "plt.legend(loc='upper left')\n",
    "plt.ylabel('Average loss')\n",
    "plt.yscale('log')\n",
    "plt.xlabel('loss Updates')\n",
    "\n",
    "plt.show()"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "averages30 = []\n",
    "averages = []\n",
    "\n",
    "for i in range(len(episodeScores)):\n",
    "    averages.append(np.array(episodeScores[max(i-100,0):i+1]).mean())\n",
    "    averages30.append(np.array(episodeScores[max(i-30,0):i+1]).mean())\n",
    "\n",
    "plt.figure(1,(10,6),1000)\n",
    "plt.plot(range(len(episodeScores)),episodeScores,alpha=0.4,label='actual',linewidth=0.5)\n",
    "plt.plot(range(len(averages)),averages,label='100 average')\n",
    "plt.plot(range(len(averages30)),averages30,label='30 average',alpha=0.6)\n",
    "plt.legend(loc='upper left')\n",
    "plt.ylabel('Average score')\n",
    "plt.xlabel('episodes')\n",
    "\n",
    "plt.show()"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  }
 ],
 "metadata": {
  "kernelspec": {
   "name": "python3",
   "language": "python",
   "display_name": "Python 3"
  },
  "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
}