{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Info\n",
    "\n",
    "Implementation of Deep Q-Network (DQN) in Tensorflow 2.0+.\n",
    "\n",
    "Uses the implementation of DQN with just the replay memory (no target memory) with bootstrap learning\n",
    "\n",
    "\n",
    "\n",
    "Resources:\n",
    "* Original DQN paper [Human-level control through deep reinforcement\n",
    "learning](https://storage.googleapis.com/deepmind-media/dqn/DQNNaturePaper.pdf).\n",
    "\n",
    "* PyTorch [REINFORCEMENT LEARNING (DQN) TUTORIAL](https://pytorch.org/tutorials/intermediate/reinforcement_q_learning.html)\n",
    "\n",
    "* [tensorflow.org RL - Intro](https://www.tensorflow.org/agents/tutorials/0_intro_rl)\n",
    "\n",
    "* [philtabor/Youtube-Code-Repository](https://github.com/philtabor/Youtube-Code-Repository)\n",
    "\n",
    "* [shivaverma/OpenAIGym](https://github.com/shivaverma/OpenAIGym)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "# from pyglet.gl import *\n",
    "# import gym\n",
    "# from IPython import display\n",
    "# import matplotlib\n",
    "# import matplotlib.pyplot as plt\n",
    "# %matplotlib inline\n",
    "\n",
    "# env = gym.make('CartPole-v0')\n",
    "# env.reset()\n",
    "# env.render(mode='rgb_array')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# try:\n",
    "#     import Image\n",
    "# except ImportError:\n",
    "#     from PIL import Image\n",
    "# import gym\n",
    "# from IPython import display\n",
    "\n",
    "# env = gym.make('Breakout-v0')\n",
    "# array = env.reset()\n",
    "# frame = PIL.Image.fromarray(env.render(mode='rgb_array'))\n",
    "# display.display(frame)\n",
    "\n",
    "# for _ in range(1000):\n",
    "#     action = env.action_space.sample()\n",
    "#     frame = PIL.Image.fromarray(env.render(mode='rgb_array'))\n",
    "#     display.display(frame)\n",
    "#     display.clear_output(wait=True)\n",
    "#     env.step(action)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'2.2.0'"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import tensorflow as tf\n",
    "import numpy as np\n",
    "import gym\n",
    "from gym import wrappers\n",
    "import matplotlib.pyplot as plt\n",
    "import pandas as pd\n",
    "plt.style.use('ggplot')\n",
    "\n",
    "# disable eager execution to speed up\n",
    "tf.compat.v1.disable_eager_execution()\n",
    "np.random.seed(0)\n",
    "\n",
    "tf.__version__"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Helper Functions"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Replay Buffer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "class ReplayBuffer(object):\n",
    "    def __init__(self, max_size, input_shape, n_actions, discrete=False):\n",
    "        \"\"\"Also called 'Replay Memory' it's designed to store previous\n",
    "            observation, action, reward, next-observation values.\n",
    "         The DQN needs experience replay to reduce correlations in \n",
    "             the sequence of observations, which otherwise might drive \n",
    "             the network into a local minimum.\n",
    "        \"\"\"\n",
    "        # when buffer gets full it goes back and re-writes again\n",
    "        self.mem_size = max_size\n",
    "        # position of first unsaved memory\n",
    "        self.mem_cntr = 0\n",
    "        self.discrete = discrete\n",
    "        self.observation_memory = np.zeros((self.mem_size, input_shape))\n",
    "        self.new_observation_memory = np.zeros((self.mem_size, input_shape))\n",
    "        dtype = np.int8 if self.discrete else np.float32\n",
    "        self.action_memory = np.zeros((self.mem_size, n_actions), dtype=dtype)\n",
    "        self.reward_memory = np.zeros(self.mem_size)\n",
    "        # keep track of our `done` flags becuase the value of the terminal state is 0.\n",
    "        self.terminal_memory = np.zeros(self.mem_size, dtype=np.float32)\n",
    "\n",
    "    def collect_experience(self, state, action, reward, state_, done):\n",
    "        \"\"\" Populate the replay memory with initial experience.\n",
    "        \"\"\"\n",
    "        index = self.mem_cntr % self.mem_size\n",
    "        self.observation_memory[index] = state\n",
    "        self.new_observation_memory[index] = state_\n",
    "        # store one hot encoding of actions, if appropriate\n",
    "        if self.discrete:\n",
    "            actions = np.zeros(self.action_memory.shape[1])\n",
    "            actions[action] = 1.0\n",
    "            self.action_memory[index] = actions\n",
    "        else:\n",
    "            self.action_memory[index] = action\n",
    "        self.reward_memory[index] = reward\n",
    "        # when episode done is 'True'. Use this to multiple the reward.\n",
    "        #  If reached temrinal state we don't give any reward (reward 0)\n",
    "        self.terminal_memory[index] = 1 - done\n",
    "        self.mem_cntr += 1\n",
    "\n",
    "    def sample_buffer(self, batch_size):\n",
    "        \"\"\"Sample memory buffer.\"\"\"\n",
    "        max_mem = min(self.mem_cntr, self.mem_size)\n",
    "        # replace will detemrine if reuse examples if already sampled\n",
    "        batch = np.random.choice(max_mem, batch_size, replace=False)\n",
    "\n",
    "        observations = self.observation_memory[batch]\n",
    "        actions = self.action_memory[batch]\n",
    "        rewards = self.reward_memory[batch]\n",
    "        new_observations = self.new_observation_memory[batch]\n",
    "        terminal = self.terminal_memory[batch]\n",
    "\n",
    "        return observations, actions, rewards, new_observations, terminal\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## DQN Agent"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "class DqnAgent(object):\n",
    "    def __init__(self, online_model, n_actions, input_dims, gamma=1, epsilon_start=0.9, epsilon_dec=0.996, epsilon_end=0.01, batch_size=32,\n",
    "                 replay_buffer_max_length=10000, fname='dqn_model.h5'):\n",
    "        \"\"\"Creat DQN Agent.\n",
    "        \n",
    "        Args:\n",
    "            online_model : TF model used for training.\n",
    "                        \n",
    "            gamma : Discount factor (0-1 value). It makes rewards from the uncertain \n",
    "                far future less important for our agent than the ones in the near \n",
    "                future that it can be fairly confident about.\n",
    "                \n",
    "            n_actions : Number of actions of environment.\n",
    "            \n",
    "            epsilon : Epsilon-greedy method for exploration.\n",
    "            \n",
    "            batch_size : Size of batch for Neural Network model training batch. Also used\n",
    "                in the memory replay, returning batch size of data.\n",
    "                \n",
    "            input_dim : Dimension of environment state.\n",
    "            \n",
    "            epsilon_dec : Decay for epislon. It starts with large value and it decays \n",
    "                to a minimum value. It performs more exploration at beginning of train.\n",
    "                \n",
    "            epsilon_end : Minimum value for epsilon. It is useful to keep a minimum value \n",
    "                of epsilon to encourage exploration durign training even after epsilon \n",
    "                finished decay.\n",
    "                \n",
    "            replay_buffer_max_length : Number of examples to store in memory that is used for 'memory replay'.\n",
    "            \n",
    "            fname : File name to save model.\n",
    "            \n",
    "        NOTE: \n",
    "            - This implementation can handle only tensorflow 2.0+.\n",
    "            - This implementaiton does not contain a target model which gets updated along training.\n",
    "        \"\"\"\n",
    "        self.action_space = list(range(n_actions))\n",
    "        self.gamma = gamma\n",
    "        self.epsilon = float(epsilon_start)\n",
    "        self.epsilon_dec = epsilon_dec\n",
    "        self.epsilon_end = epsilon_end\n",
    "        self.batch_size = batch_size\n",
    "        self.model_file = fname\n",
    "        self.memory = ReplayBuffer(replay_buffer_max_length, input_dims, n_actions,\n",
    "                                   discrete=True)\n",
    "        # online model used for training\n",
    "        self.online_model = online_model\n",
    "        # target model used as policy\n",
    "        self.target_model = tf.keras.models.clone_model(online_model)\n",
    "        \n",
    "    def collect_experience(self, observation, actions, rewards, new_observations, terminal):\n",
    "        self.memory.collect_experience(observation, actions, rewards, new_observations, terminal)\n",
    "        return\n",
    "\n",
    "    def choose_action(self, observation):\n",
    "        observation = observation[np.newaxis, :]\n",
    "        random_transition = np.random.uniform()\n",
    "        if random_transition < self.epsilon:\n",
    "            # perform exploration (random choice action)\n",
    "            action = np.random.choice(self.action_space)\n",
    "        else:\n",
    "            # perform exploitation (use current policy)\n",
    "            actions = self.target_model.predict(observation)\n",
    "            # return action value\n",
    "            action = np.argmax(actions)\n",
    "\n",
    "        return action\n",
    "    \n",
    "    def policy(self, observation):\n",
    "        \"\"\"Policy should handle observations and return best action.\n",
    "        Make sure observation is input of batch 1.\n",
    "        \"\"\"\n",
    "        \n",
    "        # expand dimension of observation if first shape is not 1\n",
    "        observation = np.expand_dims(observation, axis=0) \\\n",
    "                        if observation.shape[0]!=1 else observation\n",
    "        # perform exploitation (use current policy)\n",
    "        actions = self.target_model.predict(observation)\n",
    "        # return action value\n",
    "        action = np.argmax(actions)\n",
    "        return action\n",
    "        \n",
    "\n",
    "    def train(self):\n",
    "        if self.memory.mem_cntr > self.batch_size:\n",
    "            # fill out at least batch size memory\n",
    "            observation, action, reward, new_observation, done = \\\n",
    "                                          self.memory.sample_buffer(self.batch_size)\n",
    "\n",
    "            action_values = np.array(self.action_space, dtype=np.int8)\n",
    "            # from one-hot to indeces\n",
    "            action_indices = np.dot(action, action_values)\n",
    "            \n",
    "            # use target model as policy\n",
    "            q_eval = self.target_model.predict(observation)\n",
    "            # make copy in order to change it later\n",
    "            q_target = q_eval.copy()\n",
    "            # get Q-values for new observation\n",
    "            q_next = self.target_model.predict(new_observation)\n",
    "            \n",
    "            batch_index = np.arange(self.batch_size, dtype=np.int32)\n",
    "            # create Q-values. No reward for ending states.\n",
    "            q_target[batch_index, action_indices] = reward + \\\n",
    "                                  self.gamma*np.max(q_next, axis=1)*done\n",
    "            # fit batch to online-model\n",
    "            _ = self.online_model.fit(observation, q_target, verbose=0)\n",
    "            \n",
    "            # update epsilon\n",
    "            self.epsilon = self.epsilon*self.epsilon_dec if self.epsilon > self.epsilon_end else self.epsilon_end\n",
    "            \n",
    "#             self.epsilon = (self.epsilon - self.epsilon_dec) if self.epsilon > \\\n",
    "#                            self.epsilon_end else self.epsilon_end\n",
    "        return\n",
    "            \n",
    "    def update_policy(self):\n",
    "        \"\"\"Update target-network.\n",
    "        \"\"\"\n",
    "        self.target_model.set_weights(self.online_model.get_weights())\n",
    "        return\n",
    "    \n",
    "    def save_model(self):\n",
    "        self.q_eval.save(self.model_file)\n",
    "\n",
    "    def load_model(self):\n",
    "        self.q_eval = tf.keras.models.load_model(self.model_file)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "def encode_observation(observation_size, observation):\n",
    "    if type(observation) != np.ndarray and observation <= observation_size:\n",
    "        # need to perform one-hot encoding\n",
    "        encoded_observation = np.zeros(observation_size)\n",
    "        encoded_observation[observation] = 1\n",
    "        return encoded_observation\n",
    "    elif len(observation) == observation_size:\n",
    "        return observation\n",
    "    else:\n",
    "        raise NameError('Observation size %s does not match environment \\\n",
    "                        observation size! %s' % (observation_size, len(observation)))\n",
    "        return\n",
    "\n",
    "def compute_avg_return(environment, observation_size, policy, num_episodes=10):\n",
    "    \"\"\"Evaluate Agent on Test environment.\n",
    "    \"\"\"\n",
    "    \n",
    "    total_return = 0.0\n",
    "    for _ in range(num_episodes):\n",
    "        finished_episode = False\n",
    "        observation = environment.reset()\n",
    "        observation = encode_observation(observation_size, observation)\n",
    "        episode_return = 0.0\n",
    "        while not finished_episode:\n",
    "            # take action\n",
    "            action = policy(observation)\n",
    "            # update trajectories\n",
    "            new_observation, reward, finished_episode, info = environment.step(action)\n",
    "            new_observation = encode_observation(observation_size, new_observation)\n",
    "            # keep track of reward\n",
    "            episode_return += reward\n",
    "            # update observation\n",
    "            observation = new_observation\n",
    "        # end of episode\n",
    "        total_return += episode_return\n",
    "    # return mean of cumulated reward from each episode\n",
    "    return total_return / num_episodes"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Hyperparameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:From /users/yby1/.local/lib/python3.7/site-packages/tensorflow/python/ops/resource_variable_ops.py:1666: calling BaseResourceVariable.__init__ (from tensorflow.python.ops.resource_variable_ops) with constraint is deprecated and will be removed in a future version.\n",
      "Instructions for updating:\n",
      "If using Keras pass *_constraint arguments to layers.\n",
      "Model: \"sequential\"\n",
      "_________________________________________________________________\n",
      "Layer (type)                 Output Shape              Param #   \n",
      "=================================================================\n",
      "dense (Dense)                (None, 32)                544       \n",
      "_________________________________________________________________\n",
      "dense_1 (Dense)              (None, 64)                2112      \n",
      "_________________________________________________________________\n",
      "dense_2 (Dense)              (None, 4)                 260       \n",
      "=================================================================\n",
      "Total params: 2,916\n",
      "Trainable params: 2,916\n",
      "Non-trainable params: 0\n",
      "_________________________________________________________________\n"
     ]
    }
   ],
   "source": [
    "learning_rate = 5e-4\n",
    "batch_size = 63\n",
    "num_iterations = 15000\n",
    "policy_updates = 100\n",
    "epsilon_start = 1\n",
    "epsilon_dec = 0.999\n",
    "epsilon_end = 0.01\n",
    "gamma = 0.99\n",
    "replay_buffer_max_length  = int(1e5)  \n",
    "num_eval_episodes = 10\n",
    "eval_interval = 200\n",
    "\n",
    "\n",
    "environment_name = 'FrozenLake-v0'#'LunarLander-v2'#\"CartPole-v0\"#\"FrozenLake-v0\"\n",
    "\n",
    "train_environment = gym.make(environment_name)\n",
    "evaluate_environment = gym.make(environment_name)\n",
    "\n",
    "train_environment.seed(0)\n",
    "evaluate_environment.seed(0)\n",
    "\n",
    "n_actions = train_environment.action_space.n\n",
    "observation_size = train_environment.observation_space.shape[0] \\\n",
    "                    if train_environment.observation_space.shape \\\n",
    "                    else train_environment.observation_space.n\n",
    "\n",
    "# create tf model\n",
    "model = tf.keras.Sequential([\n",
    "        tf.keras.layers.Dense(32, input_shape=(observation_size,), activation='relu'),\n",
    "        tf.keras.layers.Dense(64, activation='relu'),\n",
    "#         tf.keras.layers.Dense(64, activation='relu'),\n",
    "        tf.keras.layers.Dense(n_actions, activation='linear')])\n",
    "\n",
    "model.compile(optimizer=tf.keras.optimizers.Adam(lr=learning_rate), loss='mse')\n",
    "\n",
    "# create DQN agent\n",
    "agent = DqnAgent(online_model=model, \n",
    "                 n_actions=n_actions, \n",
    "                 input_dims=observation_size,\n",
    "                 gamma=gamma, \n",
    "                 epsilon_start=epsilon_start, \n",
    "                 epsilon_dec=epsilon_dec, \n",
    "                 epsilon_end=epsilon_end, \n",
    "                 batch_size=batch_size, \n",
    "                 replay_buffer_max_length=replay_buffer_max_length, \n",
    "                 fname=None)\n",
    "\n",
    "# show tf model structure\n",
    "model.summary()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Train"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Iteration 200 Episodes 26 Average Return 0.10 Epsilon 0.872\n",
      "Iteration 400 Episodes 47 Average Return 0.00 Epsilon 0.714\n",
      "Iteration 600 Episodes 67 Average Return 0.20 Epsilon 0.584\n",
      "Iteration 800 Episodes 89 Average Return 0.10 Epsilon 0.478\n",
      "Iteration 1000 Episodes 112 Average Return 0.20 Epsilon 0.392\n",
      "Iteration 1200 Episodes 132 Average Return 0.10 Epsilon 0.321\n",
      "Iteration 1400 Episodes 152 Average Return 0.20 Epsilon 0.262\n",
      "Iteration 1600 Episodes 164 Average Return 0.20 Epsilon 0.215\n",
      "Iteration 1800 Episodes 178 Average Return 0.00 Epsilon 0.176\n",
      "Iteration 2000 Episodes 186 Average Return 0.50 Epsilon 0.144\n",
      "Iteration 2200 Episodes 193 Average Return 0.30 Epsilon 0.118\n",
      "Iteration 2400 Episodes 205 Average Return 0.30 Epsilon 0.097\n",
      "Iteration 2600 Episodes 212 Average Return 0.40 Epsilon 0.079\n",
      "Iteration 2800 Episodes 219 Average Return 0.40 Epsilon 0.065\n",
      "Iteration 3000 Episodes 222 Average Return 0.30 Epsilon 0.053\n",
      "Iteration 3200 Episodes 231 Average Return 0.60 Epsilon 0.043\n",
      "Iteration 3400 Episodes 235 Average Return 0.30 Epsilon 0.035\n",
      "Iteration 3600 Episodes 241 Average Return 0.70 Epsilon 0.029\n",
      "Iteration 3800 Episodes 246 Average Return 0.80 Epsilon 0.024\n",
      "Iteration 4000 Episodes 251 Average Return 0.30 Epsilon 0.019\n",
      "Iteration 4200 Episodes 255 Average Return 0.40 Epsilon 0.016\n",
      "Iteration 4400 Episodes 263 Average Return 0.50 Epsilon 0.013\n",
      "Iteration 4600 Episodes 266 Average Return 0.60 Epsilon 0.011\n",
      "Iteration 4800 Episodes 271 Average Return 0.50 Epsilon 0.010\n",
      "Iteration 5000 Episodes 279 Average Return 0.50 Epsilon 0.010\n",
      "Iteration 5200 Episodes 287 Average Return 0.80 Epsilon 0.010\n",
      "Iteration 5400 Episodes 292 Average Return 0.50 Epsilon 0.010\n",
      "Iteration 5600 Episodes 295 Average Return 0.80 Epsilon 0.010\n",
      "Iteration 5800 Episodes 304 Average Return 0.70 Epsilon 0.010\n",
      "Iteration 6000 Episodes 308 Average Return 0.60 Epsilon 0.010\n",
      "Iteration 6200 Episodes 311 Average Return 0.60 Epsilon 0.010\n",
      "Iteration 6400 Episodes 317 Average Return 1.00 Epsilon 0.010\n",
      "Iteration 6600 Episodes 320 Average Return 0.70 Epsilon 0.010\n",
      "Iteration 6800 Episodes 328 Average Return 0.60 Epsilon 0.010\n",
      "Iteration 7000 Episodes 331 Average Return 0.70 Epsilon 0.010\n",
      "Iteration 7200 Episodes 338 Average Return 0.60 Epsilon 0.010\n",
      "Iteration 7400 Episodes 341 Average Return 0.70 Epsilon 0.010\n",
      "Iteration 7600 Episodes 344 Average Return 0.90 Epsilon 0.010\n",
      "Iteration 7800 Episodes 348 Average Return 0.80 Epsilon 0.010\n",
      "Iteration 8000 Episodes 352 Average Return 0.60 Epsilon 0.010\n",
      "Iteration 8200 Episodes 357 Average Return 0.60 Epsilon 0.010\n",
      "Iteration 8400 Episodes 362 Average Return 0.80 Epsilon 0.010\n",
      "Iteration 8600 Episodes 366 Average Return 0.80 Epsilon 0.010\n",
      "Iteration 8800 Episodes 372 Average Return 0.60 Epsilon 0.010\n",
      "Iteration 9000 Episodes 374 Average Return 0.80 Epsilon 0.010\n",
      "Iteration 9200 Episodes 380 Average Return 0.70 Epsilon 0.010\n",
      "Iteration 9400 Episodes 385 Average Return 0.60 Epsilon 0.010\n",
      "Iteration 9600 Episodes 391 Average Return 0.50 Epsilon 0.010\n",
      "Iteration 9800 Episodes 398 Average Return 0.90 Epsilon 0.010\n",
      "Iteration 10000 Episodes 405 Average Return 0.60 Epsilon 0.010\n",
      "Iteration 10200 Episodes 408 Average Return 0.70 Epsilon 0.010\n",
      "Iteration 10400 Episodes 414 Average Return 0.70 Epsilon 0.010\n",
      "Iteration 10600 Episodes 421 Average Return 0.60 Epsilon 0.010\n",
      "Iteration 10800 Episodes 426 Average Return 0.50 Epsilon 0.010\n",
      "Iteration 11000 Episodes 430 Average Return 0.80 Epsilon 0.010\n",
      "Iteration 11200 Episodes 436 Average Return 0.80 Epsilon 0.010\n",
      "Iteration 11400 Episodes 443 Average Return 0.70 Epsilon 0.010\n",
      "Iteration 11600 Episodes 447 Average Return 0.60 Epsilon 0.010\n",
      "Iteration 11800 Episodes 452 Average Return 0.70 Epsilon 0.010\n",
      "Iteration 12000 Episodes 455 Average Return 0.90 Epsilon 0.010\n",
      "Iteration 12200 Episodes 460 Average Return 0.90 Epsilon 0.010\n",
      "Iteration 12400 Episodes 470 Average Return 0.60 Epsilon 0.010\n",
      "Iteration 12600 Episodes 474 Average Return 0.80 Epsilon 0.010\n",
      "Iteration 12800 Episodes 480 Average Return 0.70 Epsilon 0.010\n",
      "Iteration 13000 Episodes 487 Average Return 0.80 Epsilon 0.010\n",
      "Iteration 13200 Episodes 492 Average Return 0.70 Epsilon 0.010\n",
      "Iteration 13400 Episodes 496 Average Return 1.00 Epsilon 0.010\n",
      "Iteration 13600 Episodes 501 Average Return 0.70 Epsilon 0.010\n",
      "Iteration 13800 Episodes 508 Average Return 0.70 Epsilon 0.010\n",
      "Iteration 14000 Episodes 512 Average Return 0.90 Epsilon 0.010\n",
      "Iteration 14200 Episodes 519 Average Return 0.70 Epsilon 0.010\n",
      "Iteration 14400 Episodes 524 Average Return 0.70 Epsilon 0.010\n",
      "Iteration 14600 Episodes 528 Average Return 0.60 Epsilon 0.010\n",
      "Iteration 14800 Episodes 532 Average Return 0.80 Epsilon 0.010\n",
      "Iteration 15000 Episodes 537 Average Return 0.60 Epsilon 0.010\n"
     ]
    }
   ],
   "source": [
    "average_return = {\"iteration\":[],\n",
    "                 \"average_return\":[],\n",
    "                 }\n",
    "epsilons = []\n",
    "iteration = 0\n",
    "\n",
    "\n",
    "# evaluate the agent's policy once before training.\n",
    "avg_return = compute_avg_return(evaluate_environment, observation_size, agent.policy, num_eval_episodes)\n",
    "average_return['iteration'].append(iteration)\n",
    "average_return['average_return'].append(avg_return)\n",
    "num_episodes = 0\n",
    "\n",
    "while iteration < num_iterations:\n",
    "    finished_episode = False\n",
    "    \n",
    "    observation = train_environment.reset()\n",
    "    observation = encode_observation(observation_size, observation)\n",
    "    \n",
    "    while (not finished_episode) and (iteration < num_iterations):\n",
    "        # take action\n",
    "        action = agent.choose_action(observation)\n",
    "        # update trajectories\n",
    "        new_observation, reward, finished_episode, info = train_environment.step(action)\n",
    "        new_observation = encode_observation(observation_size, new_observation)\n",
    "        # colelct experience\n",
    "        agent.collect_experience(observation, action, reward, new_observation, int(finished_episode))\n",
    "        # update observation\n",
    "        observation = new_observation\n",
    "        # train on-line model\n",
    "        agent.train()\n",
    "        \n",
    "        # update policy\n",
    "        if iteration % policy_updates == 0:\n",
    "            agent.update_policy()\n",
    "\n",
    "        # evaluate model\n",
    "        if (iteration+1) % eval_interval  == 0:\n",
    "            avg_return = compute_avg_return(evaluate_environment, observation_size, agent.policy, num_eval_episodes)\n",
    "            \n",
    "            average_return['iteration'].append(iteration)\n",
    "            average_return['average_return'].append(avg_return)\n",
    "            print('Iteration {} Episodes {} Average Return {:.2f} Epsilon {:.3f}'.format(iteration+1,\n",
    "                                                                                         num_episodes,\n",
    "                                                                        avg_return, \n",
    "                                                                        agent.epsilon))\n",
    "        # keep track of epsilon\n",
    "        epsilons.append(agent.epsilon)\n",
    "        # next iteration\n",
    "        iteration += 1\n",
    "        \n",
    "    # end of episode\n",
    "    num_episodes += 1"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Plot"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 34788 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 18208 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "plt.figure()\n",
    "plt.title('Evaluate RL')\n",
    "plt.plot(average_return['iteration'], average_return['average_return'], label=\"Average episode reward\")\n",
    "plt.xlabel('Iterations')\n",
    "plt.ylabel('Average Return')\n",
    "plt.legend()\n",
    "\n",
    "plt.figure()\n",
    "plt.title('Epsilon Decay')\n",
    "plt.plot(epsilons, label=\"epsilon decay.\")\n",
    "plt.xlabel('Iterations')\n",
    "plt.ylabel('Epsilon')\n",
    "plt.legend()\n",
    "plt.show()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
