08a_Basics_of_Reinforcement_Learning.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## The Origins of Reinforcement Learning\n",
    "\n",
    "When it comes to analyzing the origins of modern Reinforcement Learning, there are three separate starting points that eventually merge to form what we know today: Optimal Control, Trial and Error through Animal Learning, and, less prevalent, Temporal-Difference Methods. First we will be starting with optimal control.\n",
    "\n",
    "Around the mid-1950s, Richard Bellman and others tackled the problem of \"optimal control\", described as minimizing a metric of a constantly changing enviroment over time. By combining the system's own state and a value function, optimized for a certain return goal, they were able to create a functional equation, one that is now known as the Bellman equation.\n",
    "\n",
    "This marks the beginning of what we now know as dynamic programming, the process of solving complex problems by breaking them down into subproblems and building upon each smaller solved one. Bellman also is credited with creating the Markovian decision process (MDPs) while Ronald Howard added on to MDPs by making the policy iteration method for them.\n",
    " \n",
    "The next major part is trial and error through animal learning, a practice, that according to American spychologist R. S. Woodworth, goes as far back as the late 1850s. One of the first few to truly recognize the concept of trial-and-error was Edward Thorndike, an American psychologist that worked extensively on comparitive psychology and the learning process. He initially stated what is now know as \"The Law of Effect\", a law that describes the correlation between reinforcing events and choosing actions. Over time, the theory was adapted to and laid the foundations for many professionals in the field, such as Pavlov and B. F. Skinner. \n",
    "\n",
    "In 1948, Alan Turing described a \"pleasure-pain system,\" which was expanded on and became the basis for the work of animal psychology and reinforcement learning. \n",
    "\n",
    "However, due to a lot of confusion in the previous decades due to people using the words reinforcement learning and other types of learning (such as perceptual and supervised) as synonyms, there was a period of silence where development in the field proved slow. Although, there were some exceptions to this trend. The terms \"reinforcement\" and \"reinforcement learning\" were actually used in scientific literature for the first time. This is also the time period where Minksky's paper \"Steps Toward Artificial Intelligence\" that talked about the problem of \"How  do  you  distribute  credit  for  success  among  the  many  decisions  that  may  have  been involved in producing it?\". Many topics in this paper are still relevant today. Some other examples are the system STeLLA by John Andreae and MENACE by Donald Michie.\n",
    "\n",
    "One person in particular who is attributed to reviving the field is Harry Klopf, who recognized that there were characteristics of \"adaptive behavior\" that were being fully ignored. The idea he proposed was the drive to reach a goal in the enviroment, to have a clear desired outcome and undesired end. Eventually, this push evolved into the official distinction between supervised and reinforcement learning.\n",
    "\n",
    "As mentioned previously, this is the third and last part regarding the origins of reinforcement learning: temporal-difference learning. The origins of this concept can be attributed to animal learning psychology, specifically in the idea of secondary reinforcers. A second reinforcer is a stimulus that has been passively associated with with a primary reinforer and thus has a similar effect. \n",
    "\n",
    "More information can be found in:\n",
    "- **Reinforcement Learning: An Introduction** \n",
    "2nd Edition Completed Draft, by Richard S. Sutton and Andrew G. Barto\n",
    "\n",
    "In 1989, Chris Watkin's thesis converged the major parts discussed before into developing Q-Learning.\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## States and Actions\n",
    "\n",
    "![alt text](images/states_actions.png)\n",
    "\n",
    "The first core concept we will cover is the understanding of what states and actions are. Reinforcement learning is a type of machine learning that is agent-oriented; it relies on its enviroment rather than a teacher to achieve its desired goal. This is similar to how humans learn, through the steps of trial and error.\n",
    "\n",
    "Let's take for example a person learning to navigate a maze. A state can compropise of any crossroad they are met with, an action is defined as a choice/direction they choose to go, and the goal (reward) is defined as them reaching the end of the maze.\n",
    "\n",
    "As the person navigates the maze, they will naturally discover that some paths are less optimal than others, while some do not ever reach the end. Ideally, over time, they would be able to navigate the most optimal path every time. And this is what we are trying to achieve.\n",
    "\n",
    "## Markov Decision Process \n",
    "\n",
    "Building on top of states and actions is the next step, a Markov Decision Process (MDP). A MDP can be simplified to a tuple containing 5 parts:\n",
    "   \n",
    "S - set of states   \n",
    "A - set of actions   \n",
    "P - probability that an action *a* at state *s* at time *t* will get to state *s + 1* at time *t + 1*   \n",
    "R - reward received after moving from state *s* to state *s + 1*   \n",
    "$\\gamma$ - discount factor that can optimize future rewards vs present rewards\n",
    "   \n",
    "Each of these play a role in determining a final \"policy\" $\\pi$; a rule that says given a state *s*, action *a* will be taken.\n",
    "\n",
    "![alt text](images/markov.png)\n",
    "\n",
    "This is the standard relationship between an Agent and the Enviroment in a MDP. An agent is the one who learns and makes decision while the enviroment is everything outside of the agent. These two variables constantly interact and feed each other data, with the enviroment supplying the agent with rewards and the agent triggering the effects of the enviroment.\n",
    "\n",
    "$$p({s',r | s,a}) \\doteq P\\{S_t = s', R_t = r | S_{s-1} = s , A_{t-1} = a\\}$$\n",
    "   \n",
    "What this equation states, is that in a *finite* MDP, there are a limited number of states, actions, and rewards. Because of this, we can discern that the random variables R and S have a probability distribution based only \n",
    "\n",
    "# Cross Entropy Method"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We will start with an example that we manipualted from the following book by M. Lapan: \n",
    "https://github.com/PacktPublishing/Deep-Reinforcement-Learning-Hands-On/blob/master/Chapter04/01_cartpole.py\n",
    "\n",
    "This is a nice introduction to RL with many pytorch examples included, similar to the one below. Here I have manipulated the example to run in jupyter, changed a few default behaviors, and added many more comments for understanding. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "import gym\n",
    "from collections import namedtuple\n",
    "import numpy as np\n",
    "\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "\n",
    "\n",
    "HIDDEN_SIZE = 128\n",
    "BATCH_SIZE = 16\n",
    "PERCENTILE = 70\n",
    "\n",
    "\n",
    "class Net(nn.Module):\n",
    "    def __init__(self, obs_size, hidden_size, n_actions):\n",
    "        super(Net, self).__init__()\n",
    "        self.net = nn.Sequential(\n",
    "            nn.Linear(obs_size, hidden_size),\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(hidden_size, n_actions)\n",
    "        )\n",
    "\n",
    "    def forward(self, x):\n",
    "        return self.net(x)\n",
    "\n",
    "#---------------------------------------\n",
    "# Note that we are going to use these Tuples like special classes, dictionaries \n",
    "Episode = namedtuple('Episode', field_names=['reward', 'steps'])\n",
    "EpisodeStep = namedtuple('EpisodeStep', field_names=['observation', 'action'])\n",
    "#---------------------------------------"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "# call this every iteration that we need to get\n",
    "# a batch of episodes. All environment interaction happend here \n",
    "def iterate_batches(env, net, batch_size):\n",
    "    # this function is called to generate training batches\n",
    "    # as discussed in lecture, the algorithm will \n",
    "    # try a number of random batches and return the rewards for each batch\n",
    "    # Once the total number of batches has been sampled, \n",
    "    #   we yield them for training (in a for loop below)\n",
    "    # env: an OpenAI Gym environment\n",
    "    # net: a neural network that can interact with the env\n",
    "    # batch_size: self explanatory\n",
    "    \n",
    "    batch = []\n",
    "    episode_reward = 0.0\n",
    "    episode_steps = []\n",
    "    obs = env.reset() # initial observations from environment\n",
    "    sm = nn.Softmax(dim=1) # use for translating to probability\n",
    "    while True:\n",
    "        # cast to tensor\n",
    "        obs_v = torch.FloatTensor([obs])\n",
    "        \n",
    "        # get network probabilities for an action\n",
    "        act_probs_v = sm(net(obs_v)) # use softmax here to convert to probability\n",
    "        act_probs = act_probs_v.data.numpy()[0]\n",
    "        \n",
    "        # generate an action based on probability from network\n",
    "        # this means that the network will take random actions\n",
    "        # thus exploring the space. As it trains, the actions \n",
    "        # should become more and more deterministic\n",
    "        action = np.random.choice(len(act_probs), p=act_probs)\n",
    "        \n",
    "        # take action in the environment and save obs, rewards, action\n",
    "        next_obs, reward, is_done, _ = env.step(action)\n",
    "        episode_reward += reward\n",
    "        episode_steps.append(EpisodeStep(observation=obs, action=action))\n",
    "        \n",
    "        if is_done:\n",
    "            # at the end of the episode, save the model and response\n",
    "            batch.append(Episode(reward=episode_reward, steps=episode_steps))\n",
    "            \n",
    "            # reset parameters for next episode\n",
    "            episode_reward = 0.0\n",
    "            episode_steps = []\n",
    "            next_obs = env.reset()\n",
    "            \n",
    "            # if we have generated enough episodes for a batch, \n",
    "            #  yield them to an iterator\n",
    "            if len(batch) == batch_size:\n",
    "                yield batch # yield to the user in a loop\n",
    "                # if we are here, it is after the user has \n",
    "                # used the above batch that we yielded, start over\n",
    "                batch = [] \n",
    "                \n",
    "        obs = next_obs\n",
    "\n",
    "\n",
    "# define a function to filter the various interactions\n",
    "# and only keep the best ones. We define a \"percentile\" to \n",
    "# define \"best\". The episodes that comparatively had the best reward.\n",
    "# we will loop through each episode, and put together a python list\n",
    "def filter_batch(batch, percentile):\n",
    "    # for each episode, get the reward\n",
    "    rewards = list(map(lambda s: s.reward, batch))\n",
    "    \n",
    "    # get value of the best rewards, based on current\n",
    "    reward_bound = np.percentile(rewards, percentile)\n",
    "    reward_mean = float(np.mean(rewards))\n",
    "\n",
    "    # for the best episodes, add actions/observations as training data\n",
    "    train_obs = []\n",
    "    train_act = []\n",
    "    for example in batch:\n",
    "        if example.reward >= reward_bound:\n",
    "            # extend data arrays with obs and desired actions\n",
    "            # extend adds elements to list from another list\n",
    "            #   Syntax: map(function, what to iterate over)\n",
    "            #   just get the actions and observations for this episode\n",
    "            train_obs.extend(map(lambda step: step.observation, example.steps))\n",
    "            train_act.extend(map(lambda step: step.action, example.steps))\n",
    "\n",
    "    # now each of the above are lists of observations and actions\n",
    "    # from 'good' neural networks above our reward bound\n",
    "    train_obs_v = torch.FloatTensor(train_obs)\n",
    "    train_act_v = torch.LongTensor(train_act)\n",
    "    \n",
    "    # return the best obs, actions, percentile of reward used, mean reward\n",
    "    return train_obs_v, train_act_v, reward_bound, reward_mean"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "# define cart pole environment\n",
    "env = gym.make(\"CartPole-v0\")\n",
    "\n",
    "# get the value of the observations and actions \n",
    "obs_size = env.observation_space.shape[0] # continuous space, 4 values\n",
    "n_actions = env.action_space.n # pull left or right [0,1]\n",
    "\n",
    "# create a neural network that takes the observation and gives an action\n",
    "net = Net(obs_size, HIDDEN_SIZE, n_actions)\n",
    "\n",
    "# CE method will use CE loss between the network and the \"best\" episodes\n",
    "objective = nn.CrossEntropyLoss()\n",
    "\n",
    "# and we need an optimizer to use \n",
    "optimizer = optim.Adam(params=net.parameters(), lr=0.01)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0: loss=0.699, reward_mean=18.2, reward_bound=18.5\n",
      "5: loss=0.651, reward_mean=40.2, reward_bound=49.0\n",
      "10: loss=0.604, reward_mean=59.8, reward_bound=72.0\n",
      "15: loss=0.574, reward_mean=74.5, reward_bound=80.0\n",
      "20: loss=0.553, reward_mean=85.9, reward_bound=94.0\n",
      "25: loss=0.536, reward_mean=104.0, reward_bound=122.0\n",
      "30: loss=0.512, reward_mean=112.5, reward_bound=138.0\n",
      "35: loss=0.489, reward_mean=142.7, reward_bound=147.0\n",
      "40: loss=0.459, reward_mean=143.2, reward_bound=171.5\n",
      "45: loss=0.462, reward_mean=186.9, reward_bound=200.0\n",
      "50: loss=0.464, reward_mean=181.4, reward_bound=200.0\n",
      "55: loss=0.473, reward_mean=199.7, reward_bound=200.0\n",
      "55: loss=0.473, reward_mean=199.7, reward_bound=200.0\n",
      "Solved!\n",
      "CPU times: user 54 s, sys: 3.63 s, total: 57.6 s\n",
      "Wall time: 16.6 s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "\n",
    "# iterator in for loop will yield runs of the network\n",
    "# some of these runs will, by chance, work better than others\n",
    "for iter_no, batch in enumerate(iterate_batches(env, net, BATCH_SIZE)):\n",
    "    \n",
    "    # from yielded batch, get the best actions/observation \n",
    "    #   and use as training data\n",
    "    obs_v, acts_v, reward_b, reward_m = filter_batch(batch, PERCENTILE) # we use Xth percentile\n",
    "    \n",
    "    # reset gradient calculations in graph\n",
    "    optimizer.zero_grad()\n",
    "    \n",
    "    # calculate loss of actual and selected actions \n",
    "    action_scores_v = net(obs_v) # get what the network does\n",
    "    # recall that: objective = nn.CrossEntropyLoss()\n",
    "    loss_v = objective(action_scores_v, acts_v) # use CE to define best action, no softmax\n",
    "    \n",
    "    # now back prop the gradient and update\n",
    "    loss_v.backward()\n",
    "    optimizer.step()\n",
    "    \n",
    "    if iter_no %5==0:\n",
    "        print(\"%d: loss=%.3f, reward_mean=%.1f, reward_bound=%.1f\" % (\n",
    "            iter_no, loss_v.item(), reward_m, reward_b))\n",
    "\n",
    "    if reward_m > 199:\n",
    "        print(\"%d: loss=%.3f, reward_mean=%.1f, reward_bound=%.1f\" % (\n",
    "            iter_no, loss_v.item(), reward_m, reward_b))\n",
    "        print(\"Solved!\")\n",
    "        break\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from IPython.display import clear_output, display\n",
    "\n",
    "#env = gym.make(\"CartPole-v0\")\n",
    "\n",
    "obs = env.reset()\n",
    "sm = nn.Softmax(dim=1)\n",
    "is_done = False\n",
    "while not is_done:\n",
    "    # convert to tensor\n",
    "    obs_v = torch.FloatTensor([obs])\n",
    "    # run through network to action probabilities\n",
    "    act_probs_v = sm(net(obs_v))\n",
    "    act_probs = act_probs_v.data.numpy()[0] # convert to numpy\n",
    "    # sample action according to probabilites\n",
    "    action = np.random.choice(len(act_probs), p=act_probs)\n",
    "    # take the action\n",
    "    obs, reward, is_done, _ = env.step(action)\n",
    "    \n",
    "    # display the cart\n",
    "    clear_output(wait=True)\n",
    "    result = env.render(mode=\"not_human\")\n",
    "    display(result)\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "env.close()# calling this will close and delete the current environment"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Using Cross Entropy on the Frozen Lake\n",
    "The setup of the lake is as follows:\n",
    "Observations space, integer, based on the square you select. There are holes, frozen spaces, and a goal. The reward only happens at the end, otherwise the reward is zero.\n",
    "```\n",
    "SFFF\n",
    "FHFH\n",
    "FFFH\n",
    "HFFG\n",
    "```\n",
    "\n",
    "To encode this observation space, we will convert the integer value (1-16) into a one hot encoded categorical value. \n",
    "\n",
    "The action space is defined as moving  left(1), right(2), up(3), down(4), which will be the output of the network."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Wrap the environment so that we can change default observation\n",
    "class DiscreteOneHotWrapper(gym.ObservationWrapper):\n",
    "    def __init__(self, env):\n",
    "        super(DiscreteOneHotWrapper, self).__init__(env)\n",
    "        assert isinstance(env.observation_space, gym.spaces.Discrete)\n",
    "        # change observation space to one hot encoded version \n",
    "        # we do this so that our neural network can stay the same\n",
    "        # this defines the vector of length N, with values of 0.0 up to 1.0\n",
    "        # In the gym a box is like a tensor (ugh)\n",
    "        self.observation_space = gym.spaces.Box(0.0, 1.0, (env.observation_space.n, ), \n",
    "                                                dtype=np.float32)\n",
    "\n",
    "    def observation(self, observation):\n",
    "        # get the observation lower bound (zeros) \n",
    "        res = np.zeros(self.observation_space.shape)\n",
    "        res[observation] = 1.0 # set the one hot value\n",
    "        return res"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "16 4\n"
     ]
    }
   ],
   "source": [
    "# Does it work on the frozen lake problem?\n",
    "env = DiscreteOneHotWrapper(gym.make(\"FrozenLake-v0\"))\n",
    "\n",
    "obs_size = env.observation_space.shape[0]\n",
    "n_actions = env.action_space.n\n",
    "\n",
    "# we can use these for configuring the network\n",
    "# net = Net(obs_size, HIDDEN_SIZE, n_actions)\n",
    "\n",
    "print(obs_size,n_actions)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\u001b[41mS\u001b[0mFFF\n",
      "FHFH\n",
      "FFFH\n",
      "HFFG\n"
     ]
    }
   ],
   "source": [
    "env.reset()\n",
    "env.render()\n",
    "# S: start, F: frozen, H: Hole, G: goal\n",
    "# Each action has 33% chance of going left of desired"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "  (Up)\n",
      "\u001b[41mS\u001b[0mFFF\n",
      "FHFH\n",
      "FFFH\n",
      "HFFG\n"
     ]
    }
   ],
   "source": [
    "env.step(3)\n",
    "env.render()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0: loss=1.371, reward_mean=0.0, reward_bound=0.0\n",
      "100: loss=1.316, reward_mean=0.1, reward_bound=0.0\n",
      "200: loss=1.321, reward_mean=0.1, reward_bound=0.0\n",
      "300: loss=1.059, reward_mean=0.2, reward_bound=1.0\n",
      "400: loss=0.846, reward_mean=0.2, reward_bound=1.0\n",
      "500: loss=0.694, reward_mean=0.2, reward_bound=1.0\n",
      "600: loss=0.690, reward_mean=0.2, reward_bound=1.0\n",
      "700: loss=0.612, reward_mean=0.3, reward_bound=1.0\n",
      "800: loss=0.582, reward_mean=0.3, reward_bound=1.0\n",
      "900: loss=0.512, reward_mean=0.3, reward_bound=1.0\n",
      "1000: loss=0.443, reward_mean=0.3, reward_bound=1.0\n",
      "1100: loss=0.473, reward_mean=0.3, reward_bound=1.0\n",
      "1200: loss=0.402, reward_mean=0.3, reward_bound=1.0\n",
      "1300: loss=0.444, reward_mean=0.3, reward_bound=1.0\n",
      "1400: loss=0.345, reward_mean=0.3, reward_bound=1.0\n",
      "1500: loss=0.257, reward_mean=0.3, reward_bound=1.0\n",
      "1600: loss=0.307, reward_mean=0.4, reward_bound=1.0\n",
      "1700: loss=0.202, reward_mean=0.4, reward_bound=1.0\n",
      "1800: loss=0.191, reward_mean=0.4, reward_bound=1.0\n",
      "1900: loss=0.191, reward_mean=0.4, reward_bound=1.0\n",
      "2000: loss=0.191, reward_mean=0.4, reward_bound=1.0\n",
      "Failed to converge, reached max number of iterations\n"
     ]
    }
   ],
   "source": [
    "# same code as before, but with a different environment\n",
    "PERCENTILE = 90\n",
    "net = Net(obs_size, HIDDEN_SIZE, n_actions)\n",
    "objective = nn.CrossEntropyLoss()\n",
    "optimizer = optim.Adam(params=net.parameters(), lr=0.001)\n",
    "BATCH_SIZE = 32\n",
    "reward_m_old = 0.0\n",
    "\n",
    "for iter_no, batch in enumerate(iterate_batches(env, net, BATCH_SIZE)):\n",
    "    \n",
    "    obs_v, acts_v, reward_b, reward_m = filter_batch(batch, PERCENTILE)\n",
    "    \n",
    "    # if the last values were better, use those\n",
    "    # TODO: could we try something new here?\n",
    "    if reward_m_old > reward_m:\n",
    "        obs_v, acts_v, reward_b, reward_m = obs_v_old, acts_v_old, reward_b_old, reward_m_old\n",
    "    \n",
    "    optimizer.zero_grad()\n",
    "    action_scores_v = net(obs_v)\n",
    "    loss_v = objective(action_scores_v, acts_v)\n",
    "    loss_v.backward()\n",
    "    optimizer.step()\n",
    "    \n",
    "    obs_v_old, acts_v_old, reward_b_old, reward_m_old = obs_v, acts_v, reward_b, reward_m\n",
    "    \n",
    "    if iter_no %100==0:\n",
    "        print(\"%d: loss=%.3f, reward_mean=%.1f, reward_bound=%.1f\" % (\n",
    "            iter_no, loss_v.item(), reward_m, reward_b))\n",
    "\n",
    "    if reward_m > 0.8:\n",
    "        print(\"Solved!\")\n",
    "        break\n",
    "        \n",
    "    if iter_no > 2000:\n",
    "        print(\"Failed to converge, reached max number of iterations\")\n",
    "        break"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "____\n",
    "\n",
    "**Why was this not working?**\n",
    "\n",
    "Firstly, the input space is sparse so its harder to learn new observations from the randomized neural network, especially for rarely occurring observations (like when we get past the first few steps). Also, the reward is only given at the end and its unlikely for us to reach the end, so we need to do alot of exploring... And most of the time, there is not percentile that actually worked, so we never learn to emulate the output. \n",
    "\n",
    "It seems like even this simple problem is hard for cross entropy to solve. Perhaps we should go back to the basics of learning optimal policies? Yes! Let's see about value iteration.\n",
    "\n",
    "**[Back to Slides]**"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "\n",
    "# Basics of Value Iteration\n",
    "\n",
    "https://github.com/Shmuma/Deep-Reinforcement-Learning-Hands-On/blob/master/Chapter05/01_frozenlake_v_iteration.py"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "import gym\n",
    "import collections\n",
    "\n",
    "ENV_NAME = \"FrozenLake-v0\"\n",
    "GAMMA = 0.9\n",
    "TEST_EPISODES = 20\n",
    "\n",
    "\n",
    "class Val_Agent:\n",
    "    def __init__(self, env):\n",
    "        self.env = env\n",
    "        # init reward, transitions, and value function\n",
    "        self.state = self.env.reset()\n",
    "        \n",
    "        # we will use dictionaries to be efficient\n",
    "        # tuple indexing into this will be key \n",
    "        \n",
    "        # r_{s,a,s_hat} (keep track of reward)\n",
    "        self.rewards = collections.defaultdict(float)\n",
    "        # p_{a,s to s_hat}, a dictionary of counters\n",
    "        self.transits = collections.defaultdict(collections.Counter)\n",
    "        # V(s)\n",
    "        self.values = collections.defaultdict(float)\n",
    "\n",
    "    def play_n_random_steps(self, count):\n",
    "        # play this and save the observed rewards and actions\n",
    "        for _ in range(count):\n",
    "            # randomly sample the space\n",
    "            # can get computational here, especially if we keep failing\n",
    "            action = self.env.action_space.sample()  \n",
    "            new_state, reward, is_done, _ = self.env.step(action)\n",
    "            \n",
    "            # track the reward from this action and states\n",
    "            self.rewards[(self.state, action, new_state)] = reward\n",
    "            \n",
    "            \n",
    "            # keep track of rewards to \n",
    "            #   estimate p_{a,s\\rightarrow s'}\n",
    "            self.transits[(self.state, action)][new_state] += 1\n",
    "            \n",
    "            # reset if the steps \n",
    "            self.state = self.env.reset() if is_done else new_state\n",
    "            \n",
    "            \n",
    "    def select_action(self, state):\n",
    "        # for each action, get Value of next state and reward, then choose the best\n",
    "        best_action, best_value = None, None\n",
    "        for action in range(self.env.action_space.n):\n",
    "            # get value of taking this actions, based on V and p\n",
    "            action_value = self.calc_action_value(state, action)\n",
    "            # if the best action, save this as the one to do\n",
    "            if best_value is None or best_value < action_value:\n",
    "                best_value = action_value\n",
    "                best_action = action\n",
    "        # return action with best expected reward        \n",
    "        return best_action\n",
    "    \n",
    "    def calc_action_value(self, state, action):\n",
    "        # get value of an action from a state\n",
    "        \n",
    "        # this calclates the value function for this action\n",
    "        target_counts = self.transits[(state, action)]\n",
    "        total = sum(target_counts.values()) # transition denominator\n",
    "        action_value = 0.0\n",
    "        # go through each possible target state and sum expected\n",
    "        # reward for this action (also incorporate transition)\n",
    "        for tgt_state, count in target_counts.items():\n",
    "            # get rewards from playing these steps\n",
    "            reward = self.rewards[(state, action, tgt_state)]\n",
    "            \n",
    "            # action=\\sum p_{a,s\\rightarrow s'}(r+\\gamma V(s'))\n",
    "            action_value += (count / total) * (reward + GAMMA * self.values[tgt_state])\n",
    "        return action_value\n",
    "\n",
    "\n",
    "    def play_episode(self, render=False):\n",
    "        total_reward = 0.0\n",
    "        state = self.env.reset()\n",
    "        while True:\n",
    "            # follow our policy based on Value\n",
    "            action = self.select_action(state)\n",
    "            new_state, reward, is_done, _ = self.env.step(action)\n",
    "            self.rewards[(state, action, new_state)] = reward\n",
    "            \n",
    "            self.transits[(state, action)][new_state] += 1\n",
    "            total_reward += reward\n",
    "            if render:\n",
    "                self.env.render()\n",
    "            if is_done:\n",
    "                break\n",
    "            state = new_state\n",
    "        return total_reward\n",
    "\n",
    "    def value_iteration(self):\n",
    "        # update all the values so that we can optimize. \n",
    "        # For each state, calc best values from all actions, set to maximum\n",
    "        # V(s)=\\max_{a}\\sum_{s_hat}p_{a,s,s_hat}(r_{s,a,s_hat}+\\gamma V(s_hat))\n",
    "        for state in range(self.env.observation_space.n):\n",
    "            # get values for each action\n",
    "            # \\sum_{s_hat}p_{a,s,s_hat}(r_{s,a,s_hat}+\\gamma V(s_hat))\n",
    "            state_values = [self.calc_action_value(state, action)\n",
    "                            for action in range(self.env.action_space.n)]\n",
    "            # get the max of this and save it for this state\n",
    "            self.values[state] = max(state_values)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Best reward updated 0.000 -> 0.050\n",
      "Best reward updated 0.050 -> 0.250\n",
      "Best reward updated 0.250 -> 0.300\n",
      "Best reward updated 0.300 -> 0.500\n",
      "Best reward updated 0.500 -> 0.600\n",
      "Best reward updated 0.600 -> 0.700\n",
      "Best reward updated 0.700 -> 0.800\n",
      "Best reward updated 0.800 -> 0.850\n",
      "Solved in 61 iterations!\n"
     ]
    }
   ],
   "source": [
    "env = gym.make(ENV_NAME)\n",
    "agent = Val_Agent(env)\n",
    "\n",
    "iter_no = 0\n",
    "best_reward = 0.0\n",
    "while True:\n",
    "    iter_no += 1\n",
    "    \n",
    "    # interact randomly\n",
    "    agent.play_n_random_steps(100)\n",
    "    # update matrices\n",
    "    agent.value_iteration()\n",
    "\n",
    "    # interact using the value function (and also update matrices)\n",
    "    reward = 0.0\n",
    "    for _ in range(TEST_EPISODES):\n",
    "        reward += agent.play_episode()\n",
    "    reward /= TEST_EPISODES\n",
    "    \n",
    "    if reward > best_reward:\n",
    "        print(\"Best reward updated %.3f -> %.3f\" % (best_reward, reward))\n",
    "        best_reward = reward\n",
    "    if reward > 0.80:\n",
    "        print(\"Solved in %d iterations!\" % iter_no)\n",
    "        break\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "  (Left)\n",
      "SFFF\n",
      "\u001b[41mF\u001b[0mHFH\n",
      "FFFH\n",
      "HFFG\n",
      "  (Left)\n",
      "SFFF\n",
      "FHFH\n",
      "\u001b[41mF\u001b[0mFFH\n",
      "HFFG\n",
      "  (Up)\n",
      "SFFF\n",
      "FHFH\n",
      "\u001b[41mF\u001b[0mFFH\n",
      "HFFG\n",
      "  (Up)\n",
      "SFFF\n",
      "FHFH\n",
      "\u001b[41mF\u001b[0mFFH\n",
      "HFFG\n",
      "  (Up)\n",
      "SFFF\n",
      "FHFH\n",
      "F\u001b[41mF\u001b[0mFH\n",
      "HFFG\n",
      "  (Down)\n",
      "SFFF\n",
      "FHFH\n",
      "\u001b[41mF\u001b[0mFFH\n",
      "HFFG\n",
      "  (Up)\n",
      "SFFF\n",
      "FHFH\n",
      "F\u001b[41mF\u001b[0mFH\n",
      "HFFG\n",
      "  (Down)\n",
      "SFFF\n",
      "FHFH\n",
      "FF\u001b[41mF\u001b[0mH\n",
      "HFFG\n",
      "  (Left)\n",
      "SFFF\n",
      "FHFH\n",
      "FFFH\n",
      "HF\u001b[41mF\u001b[0mG\n",
      "  (Down)\n",
      "SFFF\n",
      "FHFH\n",
      "FFFH\n",
      "HF\u001b[41mF\u001b[0mG\n",
      "  (Down)\n",
      "SFFF\n",
      "FHFH\n",
      "FFFH\n",
      "HF\u001b[41mF\u001b[0mG\n",
      "  (Down)\n",
      "SFFF\n",
      "FHFH\n",
      "FFFH\n",
      "HFF\u001b[41mG\u001b[0m\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "1.0"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "agent.play_episode(render=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "\n",
    "\n",
    "# Basics of Value Iteration with Q-Function\n",
    "Here we will simply replace the $V(s)$ with $Q(s,a)$ and update the functions. Probability will also still be kept. \n",
    "\n",
    "This means we need to update the action selection to use $Q$, and we need to update the iteration function to update $Q(s,a)$, now $V(s)$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [],
   "source": [
    "class QAgent(Val_Agent):\n",
    "\n",
    "    \n",
    "    # We inherit and use the following:\n",
    "    #     self.__init__()\n",
    "    #     self.play_episode()\n",
    "    #     self.play_n_random_steps()\n",
    "    \n",
    "    # We will abuse notation and save the Q table as self.value\n",
    "    \n",
    "    # here self.value will be Q, given by (state,action) rather than state\n",
    "    def select_action(self, state):\n",
    "        # select the best action via our Q-Function\n",
    "        best_action, best_value = None, None\n",
    "        # for each next possible action, get the best one value\n",
    "        #  this function returns a' for V(s')=max_a Q(s',a')\n",
    "        for action in range(self.env.action_space.n):\n",
    "            # get value for each action here\n",
    "            # and save the best one\n",
    "            action_value = self.values[(state, action)]\n",
    "            if best_value is None or best_value < action_value:\n",
    "                best_value = action_value\n",
    "                best_action = action\n",
    "        return best_action\n",
    "    \n",
    "    def value_iteration(self):\n",
    "        # this value is calculated from our Q function\n",
    "        # action=\\sum p_{a,s\\rightarrow s'}(r+\\gamma V(s'))\n",
    "        #   where V(s')=max_a Q(s',a') is calced from select_action\n",
    "        for state in range(self.env.observation_space.n):\n",
    "            for action in range(self.env.action_space.n):\n",
    "                action_value = 0.0\n",
    "                target_counts = self.transits[(state, action)]\n",
    "                total = sum(target_counts.values())\n",
    "                \n",
    "                for tgt_state, count in target_counts.items():\n",
    "                    reward = self.rewards[(state, action, tgt_state)]\n",
    "                    best_action = self.select_action(tgt_state)\n",
    "                    # now use best action to get V(s')\n",
    "                    #  \\sum        p_{a,s\\rightarrow s'}(    r   +\\gamma    \\max Q(s',a_best))\n",
    "                    action_value += (count / total)     * (reward + GAMMA * self.values[(tgt_state, best_action)])\n",
    "                self.values[(state, action)] = action_value"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Best reward updated 0.000 -> 0.350\n",
      "Best reward updated 0.350 -> 0.400\n",
      "Best reward updated 0.400 -> 0.450\n",
      "Best reward updated 0.450 -> 0.650\n",
      "Best reward updated 0.650 -> 0.800\n",
      "Best reward updated 0.800 -> 0.850\n",
      "Solved in 22 iterations!\n"
     ]
    }
   ],
   "source": [
    "agent = QAgent(env)\n",
    "\n",
    "iter_no = 0\n",
    "best_reward = 0.0\n",
    "while True:\n",
    "    iter_no += 1\n",
    "    \n",
    "    # same code as above, but with Q variant\n",
    "    agent.play_n_random_steps(100)\n",
    "    agent.value_iteration()\n",
    "\n",
    "    reward = 0.0\n",
    "    for _ in range(TEST_EPISODES):\n",
    "        reward += agent.play_episode()\n",
    "    reward /= TEST_EPISODES\n",
    "    if reward > best_reward:\n",
    "        print(\"Best reward updated %.3f -> %.3f\" % (best_reward, reward))\n",
    "        best_reward = reward\n",
    "    if reward > 0.8:\n",
    "        print(\"Solved in %d iterations!\" % iter_no)\n",
    "        break\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "  (Left)\n",
      "SFFF\n",
      "\u001b[41mF\u001b[0mHFH\n",
      "FFFH\n",
      "HFFG\n",
      "  (Left)\n",
      "SFFF\n",
      "FHFH\n",
      "\u001b[41mF\u001b[0mFFH\n",
      "HFFG\n",
      "  (Up)\n",
      "SFFF\n",
      "FHFH\n",
      "\u001b[41mF\u001b[0mFFH\n",
      "HFFG\n",
      "  (Up)\n",
      "SFFF\n",
      "FHFH\n",
      "\u001b[41mF\u001b[0mFFH\n",
      "HFFG\n",
      "  (Up)\n",
      "SFFF\n",
      "\u001b[41mF\u001b[0mHFH\n",
      "FFFH\n",
      "HFFG\n",
      "  (Left)\n",
      "SFFF\n",
      "FHFH\n",
      "\u001b[41mF\u001b[0mFFH\n",
      "HFFG\n",
      "  (Up)\n",
      "SFFF\n",
      "FHFH\n",
      "\u001b[41mF\u001b[0mFFH\n",
      "HFFG\n",
      "  (Up)\n",
      "SFFF\n",
      "FHFH\n",
      "F\u001b[41mF\u001b[0mFH\n",
      "HFFG\n",
      "  (Down)\n",
      "SFFF\n",
      "FHFH\n",
      "FFFH\n",
      "H\u001b[41mF\u001b[0mFG\n",
      "  (Right)\n",
      "SFFF\n",
      "FHFH\n",
      "FFFH\n",
      "HF\u001b[41mF\u001b[0mG\n",
      "  (Down)\n",
      "SFFF\n",
      "FHFH\n",
      "FFFH\n",
      "HF\u001b[41mF\u001b[0mG\n",
      "  (Down)\n",
      "SFFF\n",
      "FHFH\n",
      "FFFH\n",
      "HFF\u001b[41mG\u001b[0m\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "1.0"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "agent.play_episode(render=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**[Back to Slides]**\n",
    "\n",
    "\n",
    "# Basics of Tabular Q-learning\n",
    "\n",
    "## Q-Learning\n",
    "\n",
    "This brings us to building our first algorithm, Q-Learning. Given a state $s$, and an action $a$, the Q function returns an estimate of the total reward starting from $s$ and taking $a$.\n",
    "\n",
    "Let's go over the formula, which no longer contains any modeling of the transition probabilities, and the summation over states is dropped for an easier computation (see slides):   \n",
    " \n",
    "$$Q({s_t, a_t}) \\leftarrow (1-\\alpha)\\cdot Q({s_t, a_t}) + \\alpha[r_{t+1} +\\gamma\\max_a Q(s_{t+1}, a)]$$\n",
    "   \n",
    "$\\alpha$ - the learning rate, typically a small value between 0 and 1, indicates how much we update over values every time we take an action. Typically this value tends to be smaller in order not to overrepresent certain action. However it can also be 1, so that the $Q(s_t, a_t)$ terms cancel out.\n",
    "    \n",
    "$\\gamma$ - discount factor, encourages an agent to seek a reward sooner than later, typically set between .9 and .99. This makes agents receive a smaller reward in the present to give better incentive for future rewards. The effect of the discount factor can be seen when the Bellman equation is expanded, and $\\alpha = 1$.\n",
    "\n",
    "$$Q({s_t, a_t})_{t=0} = r_0 + \\gamma r_1 + \\gamma^2 r_2 + \\gamma^3 r_ 3 ... $$\n",
    "$$Q({s_t, a_t})_{t=0} = r_0 + \\gamma(r_1 + \\gamma^2 r_2 + \\gamma^3 r_ 3 ...) $$\n",
    "$$Q({s_t, a_t})_{t=0} = r_0 + \\gamma\\max_\\hat{a} Q(s_{1}, \\hat{a} )$$\n",
    "\n",
    "More generically, we can write this as an approximation here:\n",
    "\n",
    "$$Q({s_t, a_t}) = r_t + \\gamma\\max_\\hat{a} Q(s_{t+1}, \\hat{a} )$$\n",
    "\n",
    "<img src=\"images/q.png\" width=\"400\">\n",
    "\n",
    "Given this formula, you need the apply it using the following steps:\n",
    "1. Set initial value of *Q(s, a)* to all arbitrary values.   \n",
    "2. Eventually while reaching the limit, make sure to do all actions *a* for all states *s*.\n",
    "3. At each time *t*, change one element.\n",
    "4. You could reduce the $\\alpha$ element over time for optimization purposes.   \n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "import gym\n",
    "import collections\n",
    "\n",
    "GAMMA = 0.9\n",
    "ALPHA = 0.2\n",
    "TEST_EPISODES = 20\n",
    "\n",
    "\n",
    "class QLearningAgent:\n",
    "    def __init__(self,env):\n",
    "        self.env = env\n",
    "        # no longer keep track of transition probability matrix\n",
    "        self.state = self.env.reset()\n",
    "        # this is the Q approximation function, Q(s,a) \n",
    "        self.values = collections.defaultdict(float)\n",
    "\n",
    "    def sample_env(self):\n",
    "        # take one step in the environment and return SARS'\n",
    "        action = self.env.action_space.sample()\n",
    "        old_state = self.state\n",
    "        new_state, reward, is_done, _ = self.env.step(action)\n",
    "        self.state = self.env.reset() if is_done else new_state\n",
    "        return (old_state, action, reward, new_state)\n",
    "\n",
    "    def best_value_and_action(self, state):\n",
    "        # Go through Q(s,a), choose best action\n",
    "        best_value, best_action = None, None\n",
    "        # find V(s) = max_a Q(s,a)\n",
    "        # this function returns Q(s',a') and a' for V(s')\n",
    "        for action in range(self.env.action_space.n):\n",
    "            action_value = self.values[(state, action)]\n",
    "            if best_value is None or best_value < action_value:\n",
    "                best_value = action_value\n",
    "                best_action = action\n",
    "        return best_value, best_action\n",
    "\n",
    "    def value_update(self, s, a, r, next_s):\n",
    "        # update from one observation\n",
    "        best_v, _ = self.best_value_and_action(next_s)\n",
    "        # Q(s,a) = (1-alpha)Q(s,a) + \\alpha(r+gamma Q(s',a'))\n",
    "        new_Q = r + GAMMA * best_v\n",
    "        old_Q = self.values[(s, a)]\n",
    "        self.values[(s, a)] = old_Q * (1-ALPHA) + new_Q * ALPHA\n",
    "\n",
    "    def play_episode(self, env, render=False):\n",
    "        \n",
    "        # play what we have learned\n",
    "        # return how well we did\n",
    "        total_reward = 0.0\n",
    "        state = env.reset()\n",
    "        while True:\n",
    "            _, action = self.best_value_and_action(state)\n",
    "            new_state, reward, is_done, _ = env.step(action)\n",
    "            total_reward += reward\n",
    "            \n",
    "            if render:\n",
    "                env.render()\n",
    "                \n",
    "            if is_done:\n",
    "                break\n",
    "            state = new_state\n",
    "        return total_reward"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Best reward updated 0.000 -> 0.300\n",
      "Best reward updated 0.300 -> 0.350\n",
      "Best reward updated 0.350 -> 0.400\n",
      "Best reward updated 0.400 -> 0.450\n",
      "Best reward updated 0.450 -> 0.500\n",
      "Best reward updated 0.500 -> 0.600\n",
      "Best reward updated 0.600 -> 0.650\n",
      "Best reward updated 0.650 -> 0.700\n",
      "Best reward updated 0.700 -> 0.750\n",
      "Best reward updated 0.750 -> 0.800\n",
      "Best reward updated 0.800 -> 0.850\n",
      "Solved in 10103 iterations!\n"
     ]
    }
   ],
   "source": [
    "ENV_NAME = \"FrozenLake-v0\"\n",
    "test_env = gym.make(ENV_NAME)\n",
    "train_env = gym.make(ENV_NAME)\n",
    "agent = QLearningAgent(train_env)\n",
    "\n",
    "iter_no = 0\n",
    "best_reward = 0.0\n",
    "while True:\n",
    "    iter_no += 1\n",
    "    # sample one step\n",
    "    s, a, r, next_s = agent.sample_env()\n",
    "    # update Q\n",
    "    agent.value_update(s, a, r, next_s)\n",
    "\n",
    "    # test how well it works\n",
    "    reward = 0.0\n",
    "    for _ in range(TEST_EPISODES):\n",
    "        reward += agent.play_episode(test_env)\n",
    "        \n",
    "    # report progress\n",
    "    reward /= TEST_EPISODES\n",
    "    if reward > best_reward:\n",
    "        print(\"Best reward updated %.3f -> %.3f\" % (best_reward, reward))\n",
    "        best_reward = reward\n",
    "        \n",
    "    if reward > 0.80:\n",
    "        print(\"Solved in %d iterations!\" % iter_no)\n",
    "        break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "  (Left)\n",
      "\u001b[41mS\u001b[0mFFF\n",
      "FHFH\n",
      "FFFH\n",
      "HFFG\n",
      "  (Left)\n",
      "\u001b[41mS\u001b[0mFFF\n",
      "FHFH\n",
      "FFFH\n",
      "HFFG\n",
      "  (Left)\n",
      "\u001b[41mS\u001b[0mFFF\n",
      "FHFH\n",
      "FFFH\n",
      "HFFG\n",
      "  (Left)\n",
      "\u001b[41mS\u001b[0mFFF\n",
      "FHFH\n",
      "FFFH\n",
      "HFFG\n",
      "  (Left)\n",
      "\u001b[41mS\u001b[0mFFF\n",
      "FHFH\n",
      "FFFH\n",
      "HFFG\n",
      "  (Left)\n",
      "\u001b[41mS\u001b[0mFFF\n",
      "FHFH\n",
      "FFFH\n",
      "HFFG\n",
      "  (Left)\n",
      "\u001b[41mS\u001b[0mFFF\n",
      "FHFH\n",
      "FFFH\n",
      "HFFG\n",
      "  (Left)\n",
      "\u001b[41mS\u001b[0mFFF\n",
      "FHFH\n",
      "FFFH\n",
      "HFFG\n",
      "  (Left)\n",
      "\u001b[41mS\u001b[0mFFF\n",
      "FHFH\n",
      "FFFH\n",
      "HFFG\n",
      "  (Left)\n",
      "SFFF\n",
      "\u001b[41mF\u001b[0mHFH\n",
      "FFFH\n",
      "HFFG\n",
      "  (Left)\n",
      "SFFF\n",
      "FHFH\n",
      "\u001b[41mF\u001b[0mFFH\n",
      "HFFG\n",
      "  (Up)\n",
      "SFFF\n",
      "FHFH\n",
      "F\u001b[41mF\u001b[0mFH\n",
      "HFFG\n",
      "  (Down)\n",
      "SFFF\n",
      "FHFH\n",
      "\u001b[41mF\u001b[0mFFH\n",
      "HFFG\n",
      "  (Up)\n",
      "SFFF\n",
      "FHFH\n",
      "\u001b[41mF\u001b[0mFFH\n",
      "HFFG\n",
      "  (Up)\n",
      "SFFF\n",
      "\u001b[41mF\u001b[0mHFH\n",
      "FFFH\n",
      "HFFG\n",
      "  (Left)\n",
      "SFFF\n",
      "FHFH\n",
      "\u001b[41mF\u001b[0mFFH\n",
      "HFFG\n",
      "  (Up)\n",
      "SFFF\n",
      "\u001b[41mF\u001b[0mHFH\n",
      "FFFH\n",
      "HFFG\n",
      "  (Left)\n",
      "SFFF\n",
      "FHFH\n",
      "\u001b[41mF\u001b[0mFFH\n",
      "HFFG\n",
      "  (Up)\n",
      "SFFF\n",
      "FHFH\n",
      "F\u001b[41mF\u001b[0mFH\n",
      "HFFG\n",
      "  (Down)\n",
      "SFFF\n",
      "FHFH\n",
      "\u001b[41mF\u001b[0mFFH\n",
      "HFFG\n",
      "  (Up)\n",
      "SFFF\n",
      "FHFH\n",
      "\u001b[41mF\u001b[0mFFH\n",
      "HFFG\n",
      "  (Up)\n",
      "SFFF\n",
      "FHFH\n",
      "F\u001b[41mF\u001b[0mFH\n",
      "HFFG\n",
      "  (Down)\n",
      "SFFF\n",
      "FHFH\n",
      "\u001b[41mF\u001b[0mFFH\n",
      "HFFG\n",
      "  (Up)\n",
      "SFFF\n",
      "FHFH\n",
      "\u001b[41mF\u001b[0mFFH\n",
      "HFFG\n",
      "  (Up)\n",
      "SFFF\n",
      "\u001b[41mF\u001b[0mHFH\n",
      "FFFH\n",
      "HFFG\n",
      "  (Left)\n",
      "\u001b[41mS\u001b[0mFFF\n",
      "FHFH\n",
      "FFFH\n",
      "HFFG\n",
      "  (Left)\n",
      "\u001b[41mS\u001b[0mFFF\n",
      "FHFH\n",
      "FFFH\n",
      "HFFG\n",
      "  (Left)\n",
      "\u001b[41mS\u001b[0mFFF\n",
      "FHFH\n",
      "FFFH\n",
      "HFFG\n",
      "  (Left)\n",
      "SFFF\n",
      "\u001b[41mF\u001b[0mHFH\n",
      "FFFH\n",
      "HFFG\n",
      "  (Left)\n",
      "SFFF\n",
      "\u001b[41mF\u001b[0mHFH\n",
      "FFFH\n",
      "HFFG\n",
      "  (Left)\n",
      "\u001b[41mS\u001b[0mFFF\n",
      "FHFH\n",
      "FFFH\n",
      "HFFG\n",
      "  (Left)\n",
      "\u001b[41mS\u001b[0mFFF\n",
      "FHFH\n",
      "FFFH\n",
      "HFFG\n",
      "  (Left)\n",
      "SFFF\n",
      "\u001b[41mF\u001b[0mHFH\n",
      "FFFH\n",
      "HFFG\n",
      "  (Left)\n",
      "SFFF\n",
      "FHFH\n",
      "\u001b[41mF\u001b[0mFFH\n",
      "HFFG\n",
      "  (Up)\n",
      "SFFF\n",
      "FHFH\n",
      "\u001b[41mF\u001b[0mFFH\n",
      "HFFG\n",
      "  (Up)\n",
      "SFFF\n",
      "\u001b[41mF\u001b[0mHFH\n",
      "FFFH\n",
      "HFFG\n",
      "  (Left)\n",
      "SFFF\n",
      "FHFH\n",
      "\u001b[41mF\u001b[0mFFH\n",
      "HFFG\n",
      "  (Up)\n",
      "SFFF\n",
      "\u001b[41mF\u001b[0mHFH\n",
      "FFFH\n",
      "HFFG\n",
      "  (Left)\n",
      "SFFF\n",
      "FHFH\n",
      "\u001b[41mF\u001b[0mFFH\n",
      "HFFG\n",
      "  (Up)\n",
      "SFFF\n",
      "FHFH\n",
      "\u001b[41mF\u001b[0mFFH\n",
      "HFFG\n",
      "  (Up)\n",
      "SFFF\n",
      "\u001b[41mF\u001b[0mHFH\n",
      "FFFH\n",
      "HFFG\n",
      "  (Left)\n",
      "SFFF\n",
      "FHFH\n",
      "\u001b[41mF\u001b[0mFFH\n",
      "HFFG\n",
      "  (Up)\n",
      "SFFF\n",
      "FHFH\n",
      "F\u001b[41mF\u001b[0mFH\n",
      "HFFG\n",
      "  (Down)\n",
      "SFFF\n",
      "FHFH\n",
      "FFFH\n",
      "H\u001b[41mF\u001b[0mFG\n",
      "  (Right)\n",
      "SFFF\n",
      "FHFH\n",
      "FFFH\n",
      "HF\u001b[41mF\u001b[0mG\n",
      "  (Up)\n",
      "SFFF\n",
      "FHFH\n",
      "FFFH\n",
      "H\u001b[41mF\u001b[0mFG\n",
      "  (Right)\n",
      "SFFF\n",
      "FHFH\n",
      "FFFH\n",
      "HF\u001b[41mF\u001b[0mG\n",
      "  (Up)\n",
      "SFFF\n",
      "FHFH\n",
      "FF\u001b[41mF\u001b[0mH\n",
      "HFFG\n",
      "  (Left)\n",
      "SFFF\n",
      "FHFH\n",
      "FFFH\n",
      "HF\u001b[41mF\u001b[0mG\n",
      "  (Up)\n",
      "SFFF\n",
      "FHFH\n",
      "FF\u001b[41mF\u001b[0mH\n",
      "HFFG\n",
      "  (Left)\n",
      "SFFF\n",
      "FHFH\n",
      "F\u001b[41mF\u001b[0mFH\n",
      "HFFG\n",
      "  (Down)\n",
      "SFFF\n",
      "FHFH\n",
      "FFFH\n",
      "H\u001b[41mF\u001b[0mFG\n",
      "  (Right)\n",
      "SFFF\n",
      "FHFH\n",
      "FFFH\n",
      "H\u001b[41mF\u001b[0mFG\n",
      "  (Right)\n",
      "SFFF\n",
      "FHFH\n",
      "FFFH\n",
      "HF\u001b[41mF\u001b[0mG\n",
      "  (Up)\n",
      "SFFF\n",
      "FHFH\n",
      "FFFH\n",
      "HFF\u001b[41mG\u001b[0m\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "1.0"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "agent.play_episode(test_env, render=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**[Back To Slides]**\n",
    "\n",
    "\n",
    "# DQN (Deep Q Network)\n",
    "\n",
    "In 2013, researchers at DeepMind presented one of the first models combining reinforcement learning with a convolutional neural network. Using a neural network, they approximated the Q function, with the state being pixels from the Atari 2600. This model was able to outperform all previous approaches on playing six of the games, and outperforms human experts on three of the games.\n",
    "\n",
    "\n",
    "## Network Training\n",
    "For Atari, frames are cropped to 84x84 regions that capture the game playing area and converted to grayscale, then 4 frames are stacked to capture movement at each step. The resulting stack of frames is used as the state at each step.\n",
    "\n",
    "The network is then trained with RMSprop on the mean squared error of $Q(s)$ computed from the network and the actual reward received.  10,000,000 frames were used to train for each game.\n",
    "\n",
    "However, we will start more simple, and see if we can solve the Frozen lake problem.\n",
    "\n",
    "## Experience Replay\n",
    "Each state, action, reward and new state (known as transitions) are saved in a \"replay memory\", and at each step, a random sample of transitions are taken to train the network with. This is known as experience replay, and has a few benefits, including greater data efficiency (each state transition is used more than once) and more efficient learning (randomly sampled states are less correlated than sequential states). This also avoids oscillation and divergence, because the current state is not entirely dependent on the model's parameters at that time. The replay memory can holds the last 1,000,000 frames.\n",
    "\n",
    "## Target Network\n",
    "The authors of DQN followed up with another technique in which 2 separate Q networks are used, one to train, and one to calculate the target value during training. Every 10,000 steps, the parameters from the trained network are copied over to the target network. This also helps to avoid oscillation and divergence.\n",
    "\n",
    "Let's perform a similar setup with a toy problem (the Frozen Lake)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "import time\n",
    "import numpy as np\n",
    "import collections\n",
    "\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "import gym\n",
    "\n",
    "GAMMA = 0.9\n",
    "\n",
    "\n",
    "Experience = collections.namedtuple('Experience', \n",
    "                        field_names=['state', 'action', 'reward', 'done', 'new_state'])\n",
    "\n",
    "\n",
    "class ExperienceBuffer:\n",
    "    def __init__(self, capacity):\n",
    "        # this collection will keep track of observed SARS'\n",
    "        self.buffer = collections.deque(maxlen=capacity)\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.buffer)\n",
    "\n",
    "    def append(self, experience):\n",
    "        # buffer is a queue, so its first in first out\n",
    "        self.buffer.append(experience)\n",
    "\n",
    "    def sample(self, batch_size):\n",
    "        # return a random sample of the experience buffer\n",
    "        # output will be a numpy array of SARS' and \"is done\"\n",
    "        indices = np.random.choice(len(self.buffer), batch_size, replace=False)\n",
    "        states, actions, rewards, dones, next_states = zip(*[self.buffer[idx] for idx in indices])\n",
    "        return np.array(states), np.array(actions), np.array(rewards, dtype=np.float32), \\\n",
    "               np.array(dones, dtype=np.uint8), np.array(next_states)\n",
    "\n",
    "\n",
    "class Agent:\n",
    "    def __init__(self, env, exp_buffer):\n",
    "        # Agent will track replay buffer and hold environment\n",
    "        self.env = env\n",
    "        self.exp_buffer = exp_buffer\n",
    "        self._reset()\n",
    "\n",
    "    def _reset(self):\n",
    "        self.state = env.reset()\n",
    "        self.total_reward = 0.0\n",
    "\n",
    "    def play_step(self, net, epsilon=0.0, device=\"cpu\"):\n",
    "        done_reward = None\n",
    "\n",
    "        # use epsilon greedy approach for explore/exploit\n",
    "        if np.random.random() < epsilon:\n",
    "            # use rand policy\n",
    "            action = env.action_space.sample()\n",
    "        else:\n",
    "            # use Net policy\n",
    "            state_a = np.array([self.state], copy=False)\n",
    "            state_v = torch.tensor(state_a).to(device)\n",
    "            # get the q values for each action, given the state\n",
    "            q_vals_v = net(state_v) \n",
    "            # get idx of best action from this vector\n",
    "            _, act_v = torch.max(q_vals_v, dim=1)\n",
    "            action = int(act_v.item()) # get int from torch tensor\n",
    "\n",
    "        # do step in the environment\n",
    "        new_state, reward, is_done, _ = self.env.step(action)\n",
    "        self.total_reward += reward\n",
    "        #new_state = new_state\n",
    "\n",
    "        # add SARS' to replay buffer\n",
    "        exp = Experience(self.state, action, reward, is_done, new_state)\n",
    "        self.exp_buffer.append(exp)\n",
    "        self.state = new_state\n",
    "        \n",
    "        if is_done:\n",
    "            done_reward = self.total_reward\n",
    "            self._reset()\n",
    "        return done_reward\n",
    "\n",
    "\n",
    "def calc_loss(batch, net, tgt_net, device=\"cpu\"):\n",
    "    # batch: set of SARS' from replay buffer\n",
    "    # net: the network we are updating\n",
    "    # tgt_net: the reference network we use for updating Q\n",
    "    \n",
    "    # get the observed SARS' from the replay buffer\n",
    "    states, actions, rewards, dones, next_states = batch\n",
    "\n",
    "    # Two networks are passed in, one we are updating\n",
    "    #  and another that is a previous version\n",
    "    #  we use the previous network to observe Q(s,a)\n",
    "    \n",
    "    # send the observed states to Net, SARS'+Done\n",
    "    states_v = torch.tensor(states).to(device)\n",
    "    next_states_v = torch.tensor(next_states).to(device)\n",
    "    actions_v = torch.tensor(actions).to(device)\n",
    "    rewards_v = torch.tensor(rewards).to(device)\n",
    "    done_mask = torch.ByteTensor(dones).to(device)\n",
    "\n",
    "    # Need: \n",
    "    #  L=[ Q(s,a)-[r_{s,a}+\\gamma \\max_{a' \\in A}Q^*(s',a')] ]^2\n",
    "    \n",
    "    # get the Network actions for given states \n",
    "    state_action_values = net(states_v).gather(1, actions_v.unsqueeze(-1)).squeeze(-1)\n",
    "    # Q(s,a)\n",
    "    \n",
    "    # and the next resulting state \n",
    "    #  but only for states that did not end in a 'done' state\n",
    "    # \\max_{a' \\in A}Q^*(s',a')\n",
    "    next_state_values = tgt_net(next_states_v).max(1)[0]\n",
    "    next_state_values[done_mask] = 0.0 # ensures these are only rewards\n",
    "    \n",
    "    # detach the calculation we just made from computation graph\n",
    "    #  we don't want to back-propagate through this calculation\n",
    "    #  because it is just observations that we want to be true\n",
    "    #  That is, we want to change the expected values output from \n",
    "    #  the net, not the observations calculation\n",
    "    next_state_values = next_state_values.detach() # because from target network\n",
    "\n",
    "    # calc the Q function behavior we want (bellman update)\n",
    "    #     r_{s,a}+\\gamma \\max_{a' \\in A} Q^*(s',a')\n",
    "    expected_state_action_values = rewards_v + next_state_values * GAMMA \n",
    "    \n",
    "    # compare what we have to what we want, will update this via back prop\n",
    "    # L=[ Q(s,a)-[r_{s,a}+\\gamma \\max_{a' \\in A}Q^*(s',a')] ]^2\n",
    "    return nn.MSELoss()(state_action_values, expected_state_action_values)\n",
    "\n",
    "\n",
    "\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# Same as wrapper code from above\n",
    "# Wrap the environment so that we can change default observation\n",
    "class DiscreteOneHotWrapper(gym.ObservationWrapper):\n",
    "    def __init__(self, env):\n",
    "        super(DiscreteOneHotWrapper, self).__init__(env)\n",
    "        assert isinstance(env.observation_space, gym.spaces.Discrete)\n",
    "        # change observation space to one hot encoded version \n",
    "        # we do this so that ourneural network can stay the same\n",
    "        # this defines the vector of length N, with values from 0.0 to 1.0\n",
    "        # In the gym a box is like a tensor (ugh)\n",
    "        self.observation_space = gym.spaces.Box(0.0, 1.0, (env.observation_space.n, ), \n",
    "                                                dtype=np.float32)\n",
    "\n",
    "    def observation(self, observation):\n",
    "        res = np.copy(self.observation_space.low)\n",
    "        res[observation] = 1.0\n",
    "        return res\n",
    "    \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "16 4\n"
     ]
    }
   ],
   "source": [
    "DEFAULT_ENV_NAME = \"FrozenLake-v0\"\n",
    "\n",
    "env = DiscreteOneHotWrapper(gym.make(DEFAULT_ENV_NAME))\n",
    "\n",
    "obs_size = env.observation_space.shape[0]\n",
    "n_actions = env.action_space.n\n",
    "\n",
    "print(obs_size,n_actions)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Net(\n",
      "  (net): Sequential(\n",
      "    (0): Linear(in_features=16, out_features=256, bias=True)\n",
      "    (1): ReLU()\n",
      "    (2): Linear(in_features=256, out_features=128, bias=True)\n",
      "    (3): ReLU()\n",
      "    (4): Linear(in_features=128, out_features=4, bias=True)\n",
      "  )\n",
      ")\n",
      "16 4\n",
      "Best mean reward updated 0.000 -> 0.077, model saved\n",
      "300: done 35 iterations, mean reward 0.029, eps 1.00\n",
      "400: done 51 iterations, mean reward 0.039, eps 1.00\n",
      "3000: done 392 iterations, mean reward 0.030, eps 0.97\n",
      "4000: done 529 iterations, mean reward 0.030, eps 0.96\n",
      "6700: done 858 iterations, mean reward 0.010, eps 0.93\n",
      "7100: done 912 iterations, mean reward 0.020, eps 0.93\n",
      "7900: done 1021 iterations, mean reward 0.000, eps 0.92\n",
      "8000: done 1037 iterations, mean reward 0.000, eps 0.92\n",
      "8200: done 1063 iterations, mean reward 0.010, eps 0.92\n",
      "8900: done 1145 iterations, mean reward 0.010, eps 0.91\n",
      "9000: done 1161 iterations, mean reward 0.000, eps 0.91\n",
      "9500: done 1227 iterations, mean reward 0.010, eps 0.91\n",
      "9800: done 1269 iterations, mean reward 0.020, eps 0.90\n",
      "10400: done 1348 iterations, mean reward 0.000, eps 0.90\n",
      "10500: done 1363 iterations, mean reward 0.020, eps 0.90\n",
      "11600: done 1497 iterations, mean reward 0.000, eps 0.88\n",
      "12100: done 1558 iterations, mean reward 0.020, eps 0.88\n",
      "12800: done 1655 iterations, mean reward 0.030, eps 0.87\n",
      "12900: done 1665 iterations, mean reward 0.030, eps 0.87\n",
      "13400: done 1724 iterations, mean reward 0.020, eps 0.87\n",
      "13800: done 1776 iterations, mean reward 0.040, eps 0.86\n",
      "14100: done 1816 iterations, mean reward 0.040, eps 0.86\n",
      "14800: done 1900 iterations, mean reward 0.030, eps 0.85\n",
      "14900: done 1912 iterations, mean reward 0.040, eps 0.85\n",
      "15100: done 1940 iterations, mean reward 0.020, eps 0.85\n",
      "15800: done 2027 iterations, mean reward 0.010, eps 0.84\n",
      "16700: done 2133 iterations, mean reward 0.040, eps 0.83\n",
      "18900: done 2380 iterations, mean reward 0.030, eps 0.81\n",
      "19200: done 2412 iterations, mean reward 0.030, eps 0.81\n",
      "22100: done 2757 iterations, mean reward 0.010, eps 0.78\n",
      "22400: done 2796 iterations, mean reward 0.010, eps 0.78\n",
      "23100: done 2877 iterations, mean reward 0.020, eps 0.77\n",
      "24400: done 3026 iterations, mean reward 0.030, eps 0.76\n",
      "25000: done 3096 iterations, mean reward 0.030, eps 0.75\n",
      "25600: done 3171 iterations, mean reward 0.040, eps 0.74\n",
      "25700: done 3181 iterations, mean reward 0.040, eps 0.74\n",
      "26200: done 3232 iterations, mean reward 0.030, eps 0.74\n",
      "26700: done 3299 iterations, mean reward 0.010, eps 0.73\n",
      "27100: done 3335 iterations, mean reward 0.000, eps 0.73\n",
      "28800: done 3518 iterations, mean reward 0.050, eps 0.71\n",
      "30700: done 3705 iterations, mean reward 0.060, eps 0.69\n",
      "31500: done 3787 iterations, mean reward 0.040, eps 0.69\n",
      "32500: done 3902 iterations, mean reward 0.020, eps 0.68\n",
      "32700: done 3922 iterations, mean reward 0.010, eps 0.67\n",
      "32800: done 3934 iterations, mean reward 0.000, eps 0.67\n",
      "33400: done 4002 iterations, mean reward 0.030, eps 0.67\n",
      "35700: done 4244 iterations, mean reward 0.040, eps 0.64\n",
      "Best mean reward updated 0.077 -> 0.080, model saved\n",
      "Best mean reward updated 0.080 -> 0.090, model saved\n",
      "37200: done 4396 iterations, mean reward 0.050, eps 0.63\n",
      "38100: done 4489 iterations, mean reward 0.030, eps 0.62\n",
      "38200: done 4502 iterations, mean reward 0.030, eps 0.62\n",
      "38300: done 4514 iterations, mean reward 0.020, eps 0.62\n",
      "39100: done 4605 iterations, mean reward 0.030, eps 0.61\n",
      "41500: done 4844 iterations, mean reward 0.060, eps 0.58\n",
      "42300: done 4930 iterations, mean reward 0.020, eps 0.58\n",
      "43700: done 5067 iterations, mean reward 0.020, eps 0.56\n",
      "45400: done 5227 iterations, mean reward 0.060, eps 0.55\n",
      "46300: done 5306 iterations, mean reward 0.050, eps 0.54\n",
      "46600: done 5331 iterations, mean reward 0.030, eps 0.53\n",
      "46700: done 5341 iterations, mean reward 0.030, eps 0.53\n",
      "46900: done 5352 iterations, mean reward 0.040, eps 0.53\n",
      "49000: done 5543 iterations, mean reward 0.040, eps 0.51\n",
      "49400: done 5579 iterations, mean reward 0.050, eps 0.51\n",
      "49500: done 5587 iterations, mean reward 0.050, eps 0.51\n",
      "Best mean reward updated 0.090 -> 0.100, model saved\n",
      "Best mean reward updated 0.100 -> 0.110, model saved\n",
      "51300: done 5743 iterations, mean reward 0.100, eps 0.49\n",
      "51900: done 5794 iterations, mean reward 0.090, eps 0.48\n",
      "54500: done 6039 iterations, mean reward 0.060, eps 0.45\n",
      "56200: done 6188 iterations, mean reward 0.040, eps 0.44\n",
      "57800: done 6324 iterations, mean reward 0.070, eps 0.42\n",
      "58300: done 6354 iterations, mean reward 0.030, eps 0.42\n",
      "59200: done 6432 iterations, mean reward 0.030, eps 0.41\n",
      "59700: done 6469 iterations, mean reward 0.040, eps 0.40\n",
      "Best mean reward updated 0.110 -> 0.120, model saved\n",
      "61800: done 6642 iterations, mean reward 0.120, eps 0.38\n",
      "62000: done 6654 iterations, mean reward 0.120, eps 0.38\n",
      "Best mean reward updated 0.120 -> 0.130, model saved\n",
      "Best mean reward updated 0.130 -> 0.140, model saved\n",
      "Best mean reward updated 0.140 -> 0.150, model saved\n",
      "62500: done 6690 iterations, mean reward 0.150, eps 0.38\n",
      "64000: done 6818 iterations, mean reward 0.120, eps 0.36\n",
      "65100: done 6886 iterations, mean reward 0.090, eps 0.35\n",
      "66000: done 6954 iterations, mean reward 0.100, eps 0.34\n",
      "67100: done 7040 iterations, mean reward 0.110, eps 0.33\n",
      "68000: done 7100 iterations, mean reward 0.140, eps 0.32\n",
      "68800: done 7145 iterations, mean reward 0.120, eps 0.31\n",
      "Best mean reward updated 0.150 -> 0.160, model saved\n",
      "72700: done 7415 iterations, mean reward 0.160, eps 0.27\n",
      "Best mean reward updated 0.160 -> 0.170, model saved\n",
      "Best mean reward updated 0.170 -> 0.180, model saved\n",
      "75000: done 7573 iterations, mean reward 0.140, eps 0.25\n",
      "75300: done 7592 iterations, mean reward 0.120, eps 0.25\n",
      "78600: done 7812 iterations, mean reward 0.130, eps 0.21\n",
      "79100: done 7842 iterations, mean reward 0.110, eps 0.21\n",
      "80100: done 7888 iterations, mean reward 0.060, eps 0.20\n",
      "80800: done 7939 iterations, mean reward 0.120, eps 0.19\n",
      "Best mean reward updated 0.180 -> 0.190, model saved\n",
      "Best mean reward updated 0.190 -> 0.200, model saved\n",
      "Best mean reward updated 0.200 -> 0.210, model saved\n",
      "Best mean reward updated 0.210 -> 0.220, model saved\n",
      "Best mean reward updated 0.220 -> 0.230, model saved\n",
      "82300: done 8036 iterations, mean reward 0.200, eps 0.18\n",
      "Best mean reward updated 0.230 -> 0.240, model saved\n",
      "Best mean reward updated 0.240 -> 0.250, model saved\n",
      "Best mean reward updated 0.250 -> 0.260, model saved\n",
      "Best mean reward updated 0.260 -> 0.270, model saved\n",
      "84300: done 8137 iterations, mean reward 0.260, eps 0.16\n",
      "Best mean reward updated 0.270 -> 0.280, model saved\n",
      "Best mean reward updated 0.280 -> 0.290, model saved\n",
      "86400: done 8265 iterations, mean reward 0.110, eps 0.14\n",
      "Best mean reward updated 0.290 -> 0.300, model saved\n",
      "Best mean reward updated 0.300 -> 0.310, model saved\n",
      "Best mean reward updated 0.310 -> 0.320, model saved\n",
      "Best mean reward updated 0.320 -> 0.330, model saved\n",
      "Best mean reward updated 0.330 -> 0.340, model saved\n",
      "92200: done 8538 iterations, mean reward 0.180, eps 0.08\n",
      "Best mean reward updated 0.340 -> 0.350, model saved\n",
      "Best mean reward updated 0.350 -> 0.360, model saved\n",
      "Best mean reward updated 0.360 -> 0.370, model saved\n",
      "Best mean reward updated 0.370 -> 0.380, model saved\n",
      "Best mean reward updated 0.380 -> 0.390, model saved\n",
      "Best mean reward updated 0.390 -> 0.400, model saved\n",
      "Best mean reward updated 0.400 -> 0.410, model saved\n",
      "Best mean reward updated 0.410 -> 0.420, model saved\n",
      "Best mean reward updated 0.420 -> 0.430, model saved\n",
      "95500: done 8664 iterations, mean reward 0.380, eps 0.05\n",
      "Best mean reward updated 0.430 -> 0.440, model saved\n",
      "Best mean reward updated 0.440 -> 0.450, model saved\n",
      "Best mean reward updated 0.450 -> 0.460, model saved\n",
      "Best mean reward updated 0.460 -> 0.470, model saved\n",
      "Best mean reward updated 0.470 -> 0.480, model saved\n",
      "Best mean reward updated 0.480 -> 0.490, model saved\n",
      "Best mean reward updated 0.490 -> 0.500, model saved\n",
      "Best mean reward updated 0.500 -> 0.510, model saved\n",
      "Best mean reward updated 0.510 -> 0.520, model saved\n",
      "Best mean reward updated 0.520 -> 0.530, model saved\n",
      "Best mean reward updated 0.530 -> 0.540, model saved\n",
      "Best mean reward updated 0.540 -> 0.550, model saved\n",
      "Best mean reward updated 0.550 -> 0.560, model saved\n",
      "Best mean reward updated 0.560 -> 0.570, model saved\n",
      "Best mean reward updated 0.570 -> 0.580, model saved\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Best mean reward updated 0.580 -> 0.590, model saved\n",
      "Best mean reward updated 0.590 -> 0.600, model saved\n",
      "Best mean reward updated 0.600 -> 0.610, model saved\n",
      "Best mean reward updated 0.610 -> 0.620, model saved\n",
      "103300: done 8920 iterations, mean reward 0.620, eps 0.00\n",
      "Best mean reward updated 0.620 -> 0.630, model saved\n",
      "Best mean reward updated 0.630 -> 0.640, model saved\n",
      "Best mean reward updated 0.640 -> 0.650, model saved\n",
      "Best mean reward updated 0.650 -> 0.660, model saved\n",
      "Best mean reward updated 0.660 -> 0.670, model saved\n",
      "Best mean reward updated 0.670 -> 0.680, model saved\n",
      "Best mean reward updated 0.680 -> 0.690, model saved\n",
      "Best mean reward updated 0.690 -> 0.700, model saved\n",
      "Best mean reward updated 0.700 -> 0.710, model saved\n",
      "Best mean reward updated 0.710 -> 0.720, model saved\n",
      "Best mean reward updated 0.720 -> 0.730, model saved\n",
      "Best mean reward updated 0.730 -> 0.740, model saved\n",
      "108400: done 9053 iterations, mean reward 0.690, eps 0.00\n",
      "112700: done 9163 iterations, mean reward 0.580, eps 0.00\n",
      "113200: done 9172 iterations, mean reward 0.590, eps 0.00\n",
      "113900: done 9191 iterations, mean reward 0.570, eps 0.00\n",
      "121200: done 9360 iterations, mean reward 0.690, eps 0.00\n",
      "123400: done 9408 iterations, mean reward 0.710, eps 0.00\n",
      "123600: done 9415 iterations, mean reward 0.710, eps 0.00\n",
      "123800: done 9420 iterations, mean reward 0.730, eps 0.00\n",
      "125300: done 9454 iterations, mean reward 0.650, eps 0.00\n",
      "126500: done 9479 iterations, mean reward 0.630, eps 0.00\n",
      "130100: done 9561 iterations, mean reward 0.730, eps 0.00\n",
      "Best mean reward updated 0.740 -> 0.750, model saved\n",
      "Best mean reward updated 0.750 -> 0.760, model saved\n",
      "Best mean reward updated 0.760 -> 0.770, model saved\n",
      "131000: done 9585 iterations, mean reward 0.770, eps 0.00\n",
      "Best mean reward updated 0.770 -> 0.780, model saved\n",
      "Best mean reward updated 0.780 -> 0.790, model saved\n",
      "Best mean reward updated 0.790 -> 0.800, model saved\n",
      "Best mean reward updated 0.800 -> 0.810, model saved\n",
      "Solved in 132361 frames!\n"
     ]
    }
   ],
   "source": [
    "class Net(nn.Module):\n",
    "    def __init__(self, obs_size, hidden_size, n_actions):\n",
    "        super(Net, self).__init__()\n",
    "        self.net = nn.Sequential(\n",
    "            nn.Linear(obs_size, hidden_size),\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(hidden_size, int(hidden_size/2)),\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(int(hidden_size/2), n_actions)\n",
    "        )\n",
    "\n",
    "    def forward(self, x):\n",
    "        return self.net(x)\n",
    "\n",
    "\n",
    "\n",
    "HIDDEN_SIZE = 256\n",
    "\n",
    "net = Net(obs_size, HIDDEN_SIZE, n_actions)\n",
    "tgt_net = Net(obs_size, HIDDEN_SIZE, n_actions)\n",
    "print(net)\n",
    "\n",
    "\n",
    "print(obs_size,n_actions)\n",
    "\n",
    "device = \"cpu\"\n",
    "\n",
    "EPSILON_DECAY_LAST_FRAME = 10**5\n",
    "EPSILON_START = 1.0\n",
    "EPSILON_FINAL = 0.0\n",
    "\n",
    "MEAN_REWARD_BOUND = 0.8\n",
    "SYNC_TARGET_FRAMES = 50\n",
    "BATCH_SIZE = 16\n",
    "REPLAY_SIZE = 500\n",
    "REPLAY_START_SIZE = 500\n",
    "LEARNING_RATE = 1e-4\n",
    "\n",
    "buffer = ExperienceBuffer(REPLAY_SIZE)\n",
    "agent = Agent(env, buffer)\n",
    "epsilon = EPSILON_START\n",
    "\n",
    "optimizer = optim.Adam(net.parameters(), lr=LEARNING_RATE)\n",
    "total_rewards = []\n",
    "frame_idx = 0\n",
    "ts_frame = 0\n",
    "ts = time.time()\n",
    "best_mean_reward = None\n",
    "\n",
    "while True:\n",
    "    # track epsilon and cool it down\n",
    "    frame_idx += 1\n",
    "    epsilon = max(EPSILON_FINAL, EPSILON_START - frame_idx / EPSILON_DECAY_LAST_FRAME)\n",
    "\n",
    "    # play step and add to experience buffer\n",
    "    # here is where we populate the buffer according to a mix of random play and \n",
    "    # using the policy\n",
    "    reward = agent.play_step(net, epsilon, device=device)\n",
    "    \n",
    "    if reward is not None: # no errors!\n",
    "        total_rewards.append(reward)\n",
    "        ts_frame = frame_idx\n",
    "        \n",
    "        # calculate progress of rewards\n",
    "        mean_reward = np.mean(total_rewards[-100:])\n",
    "        if frame_idx % 100==0:\n",
    "            print(\"%d: done %d iterations, mean reward %.3f, eps %.2f\" % (\n",
    "                frame_idx, len(total_rewards), mean_reward, epsilon\n",
    "            ))\n",
    "            \n",
    "        # save best model thus far\n",
    "        if best_mean_reward is None or best_mean_reward < mean_reward:\n",
    "            torch.save(net.state_dict(), \"models/model-best.dat\")\n",
    "            if best_mean_reward is not None:\n",
    "                print(\"Best mean reward updated %.3f -> %.3f, model saved\" % (best_mean_reward, mean_reward))\n",
    "            best_mean_reward = mean_reward\n",
    "            \n",
    "        # quit if we have solved the problem\n",
    "        if mean_reward > 0.8:\n",
    "            print(\"Solved in %d frames!\" % frame_idx)\n",
    "            break\n",
    "\n",
    "    # check to see if Agent has played enough rounds\n",
    "    if len(buffer) < REPLAY_START_SIZE:\n",
    "        continue\n",
    "\n",
    "    # sync the networks every so often\n",
    "    if frame_idx % SYNC_TARGET_FRAMES == 0:\n",
    "        # use current state dictionary of values to overwrite tgt_net\n",
    "        tgt_net.load_state_dict(net.state_dict())\n",
    "\n",
    "    # use experience buffer and two networks to get loss \n",
    "    optimizer.zero_grad()\n",
    "    batch = buffer.sample(BATCH_SIZE) # grab some examples from buffer\n",
    "    loss_t = calc_loss(batch, net, tgt_net, device=device)\n",
    "    loss_t.backward()\n",
    "    optimizer.step()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# DQNs with Larger State Spaces\n",
    "\n",
    "Well it does seem like the Deep Q Network learned, but it required many more hyper parameter tunings than our previous work and seemed very brittle with respect to the epsilon parameter and replay buffer size. \n",
    "\n",
    "In fact, these are all significant downsides to the use of the Deep-Q network. For small state spaces, DQNs are not terribly advantageous. When the state space becomes intractible, however, is when they really shine---like when the state space is continuous."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# env = gym.make(\"CartPole-v0\")\n",
    "\n",
    "# obs_size = env.observation_space.shape[0]\n",
    "# n_actions = env.action_space.n\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "# class Net(nn.Module):\n",
    "#     def __init__(self, obs_size, hidden_size, n_actions):\n",
    "#         super(Net, self).__init__()\n",
    "#         self.net = nn.Sequential(\n",
    "#             nn.Linear(obs_size, hidden_size),\n",
    "#             nn.ReLU(),\n",
    "#             nn.Linear(hidden_size, int(hidden_size/2)),\n",
    "#             nn.ReLU(),\n",
    "#             nn.Linear(int(hidden_size/2), n_actions)\n",
    "#         )\n",
    "\n",
    "#     def forward(self, x):\n",
    "#         return self.net(x)\n",
    "\n",
    "\n",
    "\n",
    "# HIDDEN_SIZE = 64\n",
    "\n",
    "# net = Net(obs_size, HIDDEN_SIZE, n_actions)\n",
    "# tgt_net = Net(obs_size, HIDDEN_SIZE, n_actions)\n",
    "# print(net)\n",
    "\n",
    "\n",
    "# print(obs_size,n_actions)\n",
    "\n",
    "# device = \"cpu\"\n",
    "\n",
    "# EPSILON_DECAY_LAST_FRAME = 10**5\n",
    "# EPSILON_START = 1.0\n",
    "# EPSILON_FINAL = 0.0\n",
    "\n",
    "# MEAN_REWARD_BOUND = 200\n",
    "# SYNC_TARGET_FRAMES = 50\n",
    "# BATCH_SIZE = 16\n",
    "# REPLAY_SIZE = 500\n",
    "# REPLAY_START_SIZE = 500\n",
    "# LEARNING_RATE = 1e-4\n",
    "\n",
    "# buffer = ExperienceBuffer(REPLAY_SIZE)\n",
    "# agent = Agent(env, buffer)\n",
    "# epsilon = EPSILON_START\n",
    "\n",
    "# optimizer = optim.Adam(net.parameters(), lr=LEARNING_RATE)\n",
    "# total_rewards = []\n",
    "# frame_idx = 0\n",
    "# ts_frame = 0\n",
    "# ts = time.time()\n",
    "# best_mean_reward = None\n",
    "\n",
    "# while True:\n",
    "#     frame_idx += 1\n",
    "#     epsilon = max(EPSILON_FINAL, EPSILON_START - frame_idx / EPSILON_DECAY_LAST_FRAME)\n",
    "\n",
    "#     reward = agent.play_step(net, epsilon, device=device)\n",
    "#     if reward is not None:\n",
    "#         total_rewards.append(reward)\n",
    "#         ts_frame = frame_idx\n",
    "        \n",
    "#         mean_reward = np.mean(total_rewards[-100:])\n",
    "#         if frame_idx % 100==0:\n",
    "#             print(\"%d: done %d iterations, mean reward %.3f, eps %.2f\" % (\n",
    "#                 frame_idx, len(total_rewards), mean_reward, epsilon\n",
    "#             ))\n",
    "        \n",
    "#         if best_mean_reward is None or best_mean_reward < mean_reward:\n",
    "#             torch.save(net.state_dict(), \"models/model-best.dat\")\n",
    "#             if best_mean_reward is not None:\n",
    "#                 print(\"Best mean reward updated %.3f -> %.3f, model saved\" % (best_mean_reward, mean_reward))\n",
    "#             best_mean_reward = mean_reward\n",
    "#         if mean_reward > 200:\n",
    "#             print(\"Solved in %d frames!\" % frame_idx)\n",
    "#             break\n",
    "\n",
    "#     if len(buffer) < REPLAY_START_SIZE:\n",
    "#         continue\n",
    "\n",
    "#     if frame_idx % SYNC_TARGET_FRAMES == 0:\n",
    "#         tgt_net.load_state_dict(net.state_dict())\n",
    "\n",
    "#     optimizer.zero_grad()\n",
    "#     batch = buffer.sample(BATCH_SIZE)\n",
    "#     loss_t = calc_loss(batch, net, tgt_net, device=device)\n",
    "#     loss_t.backward()\n",
    "#     optimizer.step()\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Deep Q Learning with Atari Games\n",
    "\n",
    "https://github.com/Shmuma/Deep-Reinforcement-Learning-Hands-On/blob/master/Chapter06/02_dqn_pong.py"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "# load up some utilities \n",
    "from q_learn_utils import make_env\n",
    "from q_learn_utils import DQN\n",
    "\n",
    "import time\n",
    "import numpy as np\n",
    "import collections\n",
    "\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "\n",
    "\n",
    "DEFAULT_ENV_NAME = \"PongNoFrameskip-v4\"\n",
    "MEAN_REWARD_BOUND = 19.5\n",
    "\n",
    "GAMMA = 0.99\n",
    "BATCH_SIZE = 32\n",
    "REPLAY_SIZE = 10000\n",
    "LEARNING_RATE = 1e-4\n",
    "SYNC_TARGET_FRAMES = 1000\n",
    "REPLAY_START_SIZE = 10000\n",
    "\n",
    "EPSILON_DECAY_LAST_FRAME = 10**5\n",
    "EPSILON_START = 1.0\n",
    "EPSILON_FINAL = 0.02\n",
    "\n",
    "\n",
    "Experience = collections.namedtuple('Experience', \n",
    "                                    field_names=['state', 'action', 'reward', 'done', 'new_state'])\n",
    "\n",
    "\n",
    "device = \"cpu\"\n",
    "\n",
    "# load our own custom environment\n",
    "env = make_env(DEFAULT_ENV_NAME)\n",
    "# this has lots of tricks in it, including:\n",
    "# 1. press fire to start game in atari\n",
    "# 2. Max pool across frames (max across four frames, keeping last two)\n",
    "# 3. Resize, gray scale, and crop atari images (get rid of score and other unneeded pixels)\n",
    "# 4. PyTorch Image conversion\n",
    "# 5. Image scaling (input 0 to 1, rather than 0-255)\n",
    "# 6. Use last four buffer of previous observations \n",
    "\n",
    "# load up a simple convolutional network\n",
    "# 3 layers of strided conv and two fc layers\n",
    "net = DQN(env.observation_space.shape, env.action_space.n).to(device)\n",
    "tgt_net = DQN(env.observation_space.shape, env.action_space.n).to(device)\n",
    "print(net)\n",
    "\n",
    "buffer = ExperienceBuffer(REPLAY_SIZE)\n",
    "agent = Agent(env, buffer) # same as previous Q-Learning agent\n",
    "epsilon = EPSILON_START\n",
    "\n",
    "optimizer = optim.Adam(net.parameters(), lr=LEARNING_RATE)\n",
    "total_rewards = []\n",
    "frame_idx = 0\n",
    "ts_frame = 0\n",
    "ts = time.time()\n",
    "best_mean_reward = None\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "# training (no resets of the Agent or training values)\n",
    "# this is the same code that is better commented above\n",
    "# only minor changes are made to save the models\n",
    "while True:\n",
    "    frame_idx += 1\n",
    "    epsilon = max(EPSILON_FINAL, EPSILON_START - frame_idx / EPSILON_DECAY_LAST_FRAME)\n",
    "\n",
    "    # use policy to interact with environment, with epsilon\n",
    "    # play until complete \n",
    "    reward = agent.play_step(net, epsilon, device=device)\n",
    "    \n",
    "    if reward is not None:\n",
    "        # housekeeping code and saving out the best model\n",
    "        total_rewards.append(reward)\n",
    "        speed = (frame_idx - ts_frame) / (time.time() - ts)\n",
    "        ts_frame = frame_idx\n",
    "        ts = time.time()\n",
    "        mean_reward = np.mean(total_rewards[-100:])\n",
    "        \n",
    "        print(\"%d: done %d games, mean reward %.3f, eps %.2f, speed %.2f f/s\" % (\n",
    "            frame_idx, len(total_rewards), mean_reward, epsilon,\n",
    "            speed\n",
    "        ))\n",
    "        \n",
    "        if best_mean_reward is None or best_mean_reward < mean_reward:\n",
    "            torch.save(net.state_dict(),\"models/\" + DEFAULT_ENV_NAME + \"-best.dat\")\n",
    "            if best_mean_reward is not None:\n",
    "                print(\"Best mean reward updated %.3f -> %.3f, model saved\" % (best_mean_reward, mean_reward))\n",
    "            best_mean_reward = mean_reward\n",
    "            \n",
    "        if mean_reward > MEAN_REWARD_BOUND:\n",
    "            print(\"Solved in %d frames!\" % frame_idx)\n",
    "            break\n",
    "\n",
    "    if len(buffer) < REPLAY_START_SIZE:\n",
    "        continue\n",
    "\n",
    "    if frame_idx % SYNC_TARGET_FRAMES == 0:\n",
    "        # periodicially update the network, from Q to Q* (target)\n",
    "        tgt_net.load_state_dict(net.state_dict())\n",
    "\n",
    "    optimizer.zero_grad()\n",
    "    # sample experience buffer for desired Q(s,a)\n",
    "    batch = buffer.sample(BATCH_SIZE) \n",
    "    # # incentivize bellman update with momentum \n",
    "    loss_t = calc_loss(batch, net, tgt_net, device=device)  \n",
    "    loss_t.backward()\n",
    "    optimizer.step()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set this to run on my wife's computer for about three days :)\n",
    "# Thanks to my Wife!!\n",
    "\n",
    "# 420406: done 235 games, mean reward 12.430, eps 0.02, speed 0.72 f/s\n",
    "# Best mean reward updated 12.220 -> 12.430, model saved\n",
    "# 422326: done 236 games, mean reward 12.630, eps 0.02, speed 5.20 f/s\n",
    "# Best mean reward updated 12.430 -> 12.630, model saved\n",
    "# 424148: done 237 games, mean reward 12.860, eps 0.02, speed 5.54 f/s\n",
    "# Best mean reward updated 12.630 -> 12.860, model saved\n",
    "# 426031: done 238 games, mean reward 13.110, eps 0.02, speed 5.52 f/s\n",
    "# Best mean reward updated 12.860 -> 13.110, model saved\n",
    "# 428313: done 239 games, mean reward 13.280, eps 0.02, speed 5.49 f/s\n",
    "# Best mean reward updated 13.110 -> 13.280, model saved\n",
    "# 430085: done 240 games, mean reward 13.540, eps 0.02, speed 5.61 f/s\n",
    "# Best mean reward updated 13.280 -> 13.540, model saved\n",
    "# 432531: done 241 games, mean reward 13.710, eps 0.02, speed 5.34 f/s\n",
    "# Best mean reward updated 13.540 -> 13.710, model saved\n",
    "# 434197: done 242 games, mean reward 13.960, eps 0.02, speed 5.68 f/s\n",
    "# Best mean reward updated 13.710 -> 13.960, model saved\n",
    "# 435898: done 243 games, mean reward 14.290, eps 0.02, speed 5.87 f/s\n",
    "# Best mean reward updated 13.960 -> 14.290, model saved\n",
    "# 438575: done 244 games, mean reward 14.360, eps 0.02, speed 5.72 f/s\n",
    "# Best mean reward updated 14.290 -> 14.360, model saved\n",
    "# 440553: done 245 games, mean reward 14.590, eps 0.02, speed 5.78 f/s\n",
    "# Best mean reward updated 14.360 -> 14.590, model saved\n",
    "# 442649: done 246 games, mean reward 14.740, eps 0.02, speed 5.81 f/s\n",
    "# Best mean reward updated 14.590 -> 14.740, model saved\n",
    "# 444772: done 247 games, mean reward 14.850, eps 0.02, speed 5.85 f/s\n",
    "# Best mean reward updated 14.740 -> 14.850, model saved\n",
    "# 446983: done 248 games, mean reward 14.970, eps 0.02, speed 5.86 f/s\n",
    "# Best mean reward updated 14.850 -> 14.970, model saved\n",
    "# 448648: done 249 games, mean reward 15.200, eps 0.02, speed 5.86 f/s\n",
    "# Best mean reward updated 14.970 -> 15.200, model saved\n",
    "# 450580: done 250 games, mean reward 15.350, eps 0.02, speed 5.86 f/s\n",
    "# Best mean reward updated 15.200 -> 15.350, model saved\n",
    "# 452333: done 251 games, mean reward 15.490, eps 0.02, speed 5.86 f/s\n",
    "# Best mean reward updated 15.350 -> 15.490, model saved\n",
    "# 454234: done 252 games, mean reward 15.630, eps 0.02, speed 5.80 f/s\n",
    "# Best mean reward updated 15.490 -> 15.630, model saved\n",
    "# 456215: done 253 games, mean reward 15.730, eps 0.02, speed 5.59 f/s\n",
    "# Best mean reward updated 15.630 -> 15.730, model saved\n",
    "# 458373: done 254 games, mean reward 15.720, eps 0.02, speed 5.57 f/s\n",
    "# 460341: done 255 games, mean reward 15.810, eps 0.02, speed 5.53 f/s\n",
    "# Best mean reward updated 15.730 -> 15.810, model saved\n",
    "# 462813: done 256 games, mean reward 15.820, eps 0.02, speed 5.60 f/s\n",
    "# Best mean reward updated 15.810 -> 15.820, model saved\n",
    "# 464478: done 257 games, mean reward 15.920, eps 0.02, speed 5.50 f/s\n",
    "# Best mean reward updated 15.820 -> 15.920, model saved\n",
    "# 466469: done 258 games, mean reward 15.910, eps 0.02, speed 5.81 f/s\n",
    "# 468349: done 259 games, mean reward 15.910, eps 0.02, speed 5.70 f/s\n",
    "# 470003: done 260 games, mean reward 15.920, eps 0.02, speed 5.85 f/s\n",
    "# 471728: done 261 games, mean reward 15.920, eps 0.02, speed 5.84 f/s\n",
    "# 473664: done 262 games, mean reward 15.910, eps 0.02, speed 5.84 f/s\n",
    "# 475491: done 263 games, mean reward 15.910, eps 0.02, speed 5.56 f/s\n",
    "# 477926: done 264 games, mean reward 15.810, eps 0.02, speed 5.60 f/s\n",
    "# 480288: done 265 games, mean reward 15.820, eps 0.02, speed 5.63 f/s\n",
    "# 482233: done 266 games, mean reward 15.960, eps 0.02, speed 5.80 f/s\n",
    "# Best mean reward updated 15.920 -> 15.960, model saved\n",
    "# 484020: done 267 games, mean reward 16.070, eps 0.02, speed 5.93 f/s\n",
    "# Best mean reward updated 15.960 -> 16.070, model saved\n",
    "# 485689: done 268 games, mean reward 16.110, eps 0.02, speed 5.88 f/s\n",
    "# Best mean reward updated 16.070 -> 16.110, model saved\n",
    "# 487387: done 269 games, mean reward 16.160, eps 0.02, speed 5.89 f/s\n",
    "# Best mean reward updated 16.110 -> 16.160, model saved\n",
    "# 489172: done 270 games, mean reward 16.220, eps 0.02, speed 5.88 f/s\n",
    "# Best mean reward updated 16.160 -> 16.220, model saved\n",
    "# 490891: done 271 games, mean reward 16.260, eps 0.02, speed 5.84 f/s\n",
    "# Best mean reward updated 16.220 -> 16.260, model saved\n",
    "# 492973: done 272 games, mean reward 16.260, eps 0.02, speed 5.79 f/s\n",
    "# 494942: done 273 games, mean reward 16.250, eps 0.02, speed 5.87 f/s\n",
    "# 496671: done 274 games, mean reward 16.270, eps 0.02, speed 5.58 f/s\n",
    "# Best mean reward updated 16.260 -> 16.270, model saved\n",
    "# 498337: done 275 games, mean reward 16.360, eps 0.02, speed 5.61 f/s\n",
    "# Best mean reward updated 16.270 -> 16.360, model saved\n",
    "# 500189: done 276 games, mean reward 16.440, eps 0.02, speed 5.71 f/s\n",
    "# Best mean reward updated 16.360 -> 16.440, model saved\n",
    "# 502074: done 277 games, mean reward 16.470, eps 0.02, speed 5.83 f/s\n",
    "# Best mean reward updated 16.440 -> 16.470, model saved\n",
    "# 503799: done 278 games, mean reward 16.540, eps 0.02, speed 5.87 f/s\n",
    "# Best mean reward updated 16.470 -> 16.540, model saved\n",
    "# 505524: done 279 games, mean reward 16.580, eps 0.02, speed 5.82 f/s\n",
    "# Best mean reward updated 16.540 -> 16.580, model saved\n",
    "# 507347: done 280 games, mean reward 16.620, eps 0.02, speed 5.82 f/s\n",
    "# Best mean reward updated 16.580 -> 16.620, model saved\n",
    "# 509706: done 281 games, mean reward 16.580, eps 0.02, speed 5.83 f/s\n",
    "# 511668: done 282 games, mean reward 16.590, eps 0.02, speed 5.88 f/s\n",
    "# 513617: done 283 games, mean reward 16.610, eps 0.02, speed 5.90 f/s\n",
    "# 515254: done 284 games, mean reward 16.660, eps 0.02, speed 5.89 f/s\n",
    "# Best mean reward updated 16.620 -> 16.660, model saved\n",
    "# 516919: done 285 games, mean reward 16.670, eps 0.02, speed 5.67 f/s\n",
    "# Best mean reward updated 16.660 -> 16.670, model saved\n",
    "# 519092: done 286 games, mean reward 16.690, eps 0.02, speed 5.60 f/s\n",
    "# Best mean reward updated 16.670 -> 16.690, model saved\n",
    "# 521198: done 287 games, mean reward 16.700, eps 0.02, speed 5.62 f/s\n",
    "# Best mean reward updated 16.690 -> 16.700, model saved\n",
    "# 523572: done 288 games, mean reward 16.610, eps 0.02, speed 5.73 f/s\n",
    "# 525237: done 289 games, mean reward 16.610, eps 0.02, speed 5.84 f/s\n",
    "# 527041: done 290 games, mean reward 16.630, eps 0.02, speed 5.82 f/s\n",
    "# 528859: done 291 games, mean reward 16.640, eps 0.02, speed 5.78 f/s\n",
    "# 530865: done 292 games, mean reward 16.630, eps 0.02, speed 5.44 f/s\n",
    "# 532944: done 293 games, mean reward 16.620, eps 0.02, speed 5.19 f/s\n",
    "# ---------------------------------------------------------------------------\n",
    "# KeyboardInterrupt                         Traceback (most recent call last)\n",
    "\n",
    "# Ended updates after the score was pretty good \n",
    "# any positive score means we won the game"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Running Trained Agent on Pong"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<All keys matched successfully>"
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import gym\n",
    "import time\n",
    "import numpy as np\n",
    "\n",
    "import torch\n",
    "\n",
    "from q_learn_utils import make_env\n",
    "from q_learn_utils import DQN\n",
    "\n",
    "import collections\n",
    "\n",
    "DEFAULT_ENV_NAME = \"PongNoFrameskip-v4\"\n",
    "RENDER = True\n",
    "\n",
    "env = make_env(DEFAULT_ENV_NAME)\n",
    "\n",
    "test_net = DQN(env.observation_space.shape, env.action_space.n)\n",
    "test_net.load_state_dict(torch.load(\"models/\" +DEFAULT_ENV_NAME + \"-best_Spring2019.dat\", map_location=lambda storage, loc: storage))\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Total reward: 20.00\n",
      "Action counts: Counter({2: 428, 5: 390, 3: 331, 4: 228, 0: 198, 1: 191})\n"
     ]
    }
   ],
   "source": [
    "state = env.reset()\n",
    "total_reward = 0.0\n",
    "c = collections.Counter()\n",
    "\n",
    "while True:\n",
    "    start_ts = time.time()\n",
    "    if RENDER:\n",
    "        env.render()\n",
    "    state_v = torch.tensor(np.array([state], copy=False))\n",
    "    q_vals = test_net(state_v).data.numpy()[0]\n",
    "    action = np.argmax(q_vals)\n",
    "    c[action] += 1\n",
    "    state, reward, done, _ = env.step(action)\n",
    "    total_reward += reward\n",
    "    if done:\n",
    "        break\n",
    "    if RENDER:# too fast without FPS limiter\n",
    "        delta = 1/30 - (time.time() - start_ts)\n",
    "        if delta > 0:\n",
    "            time.sleep(delta)\n",
    "print(\"Total reward: %.2f\" % total_reward)\n",
    "print(\"Action counts:\", c)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "env.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}