week04_approx_rl/homework_pytorch_debug.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "SqZ2EwnTZdC8"
   },
   "source": [
    "# Deep Q-Network implementation.\n",
    "\n",
    "This homework shamelessly demands you to implement DQN — an approximate Q-learning algorithm with experience replay and target networks — and see if it works any better this way.\n",
    "\n",
    "Original paper:\n",
    "https://arxiv.org/pdf/1312.5602.pdf"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Zv7XJfXaZdC9"
   },
   "source": [
    "**This notebook is given for debug.** The main task is in the other notebook (**homework_pytorch_main**). The tasks are similar and share most of the code. The main difference is in environments. In main notebook it can take some 2 hours for the agent to start improving so it seems reasonable to launch the algorithm on a simpler env first. Here it is CartPole and it will train in several minutes.\n",
    "\n",
    "**We suggest the following pipeline:** First implement debug notebook then implement the main one.\n",
    "\n",
    "**About evaluation:** All points are given for the main notebook with one exception: if agent fails to beat the threshold in main notebook you can get 1 pt (instead of 3 pts) for beating the threshold in debug notebook."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "ioIEVODJZdC9"
   },
   "outputs": [],
   "source": [
    "import sys, os\n",
    "if 'google.colab' in sys.modules and not os.path.exists('.setup_complete'):\n",
    "    os.makedirs('dqn', exist_ok=True)\n",
    "    !wget -q https://raw.githubusercontent.com/yandexdataschool/Practical_RL/master/setup_colab.sh -O- | bash\n",
    "\n",
    "    !wget -q https://raw.githubusercontent.com/yandexdataschool/Practical_RL/master/week04_approx_rl/dqn/atari_wrappers.py -P dqn/\n",
    "    !wget -q https://raw.githubusercontent.com/yandexdataschool/Practical_RL/master/week04_approx_rl/dqn/utils.py -P dqn/\n",
    "    !wget -q https://raw.githubusercontent.com/yandexdataschool/Practical_RL/master/week04_approx_rl/dqn/replay_buffer.py -P dqn/\n",
    "    !wget -q https://raw.githubusercontent.com/yandexdataschool/Practical_RL/master/week04_approx_rl/dqn/framebuffer.py -P dqn/\n",
    "    !wget -q https://raw.githubusercontent.com/yandexdataschool/Practical_RL/master/week04_approx_rl/dqn/analysis.py -P dqn/\n",
    "\n",
    "    !pip install gymnasium\n",
    "\n",
    "    !touch .setup_complete\n",
    "\n",
    "# This code creates a virtual display to draw game images on.\n",
    "# It will have no effect if your machine has a monitor.\n",
    "if type(os.environ.get(\"DISPLAY\")) is not str or len(os.environ.get(\"DISPLAY\")) == 0:\n",
    "    !bash ../xvfb start\n",
    "    os.environ['DISPLAY'] = ':1'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "FDZqlI3kZdC9"
   },
   "source": [
    "__Frameworks__ - we'll accept this homework in any deep learning framework. This particular notebook was designed for PyTorch, but you find it easy to adapt it to almost any Python-based deep learning framework."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "dsYq558wZdC-"
   },
   "outputs": [],
   "source": [
    "import random\n",
    "import numpy as np\n",
    "import torch\n",
    "import dqn.utils\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "6ypPZ8e6ZdC-"
   },
   "outputs": [],
   "source": [
    "import gymnasium as gym\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "9j8EGNlSZdC-"
   },
   "source": [
    "### CartPole again\n",
    "\n",
    "Another env can be used without any modification of the code. State space should be a single vector, actions should be discrete.\n",
    "\n",
    "CartPole is the simplest one. It should take several minutes to solve it.\n",
    "\n",
    "For LunarLander it can take 1-2 hours to get 200 points (a good score) on Colab and training progress does not look informative."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "v-5u-CcQZdC-"
   },
   "outputs": [],
   "source": [
    "ENV_NAME = \"CartPole-v1\"\n",
    "\n",
    "\n",
    "def make_env():\n",
    "    # some envs are wrapped with a time limit wrapper by default\n",
    "    env = gym.make(ENV_NAME, render_mode=\"rgb_array\").unwrapped\n",
    "    return env\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "AmFXRrkqZdC-"
   },
   "outputs": [],
   "source": [
    "env = make_env()\n",
    "env.reset()\n",
    "plt.imshow(env.render())\n",
    "state_shape, n_actions = env.observation_space.shape, env.action_space.n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "qOyWgOmvZdC-"
   },
   "source": [
    "### Building a network"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "XqpThLZXZdC-"
   },
   "source": [
    "We now need to build a neural network that can map observations to state q-values.\n",
    "The model does not have to be huge yet. 1-2 hidden layers with < 200 neurons and ReLU activation will probably be enough. Batch normalization and dropout can spoil everything here."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "UVlpkvZOZdC-"
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
    "# those who have a GPU but feel unfair to use it can uncomment:\n",
    "# device = torch.device('cpu')\n",
    "device"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "RFva1cpyZdC-"
   },
   "outputs": [],
   "source": [
    "class DQNAgent(nn.Module):\n",
    "    def __init__(self, state_shape, n_actions, epsilon=0):\n",
    "\n",
    "        super().__init__()\n",
    "        self.epsilon = epsilon\n",
    "        self.n_actions = n_actions\n",
    "        self.state_shape = state_shape\n",
    "        # Define your network body here. Please make sure agent is fully contained here\n",
    "        assert len(state_shape) == 1\n",
    "        state_dim = state_shape[0]\n",
    "        <YOUR CODE>\n",
    "\n",
    "\n",
    "    def forward(self, state_t):\n",
    "        \"\"\"\n",
    "        takes agent's observation (tensor), returns qvalues (tensor)\n",
    "        :param state_t: a batch states, shape = [batch_size, *state_dim=4]\n",
    "        \"\"\"\n",
    "        # Use your network to compute qvalues for given state\n",
    "        qvalues = <YOUR CODE>\n",
    "\n",
    "        assert qvalues.requires_grad, \"qvalues must be a torch tensor with grad\"\n",
    "        assert (\n",
    "            len(qvalues.shape) == 2 and\n",
    "            qvalues.shape[0] == state_t.shape[0] and\n",
    "            qvalues.shape[1] == n_actions\n",
    "        )\n",
    "\n",
    "        return qvalues\n",
    "\n",
    "    def get_qvalues(self, states):\n",
    "        \"\"\"\n",
    "        like forward, but works on numpy arrays, not tensors\n",
    "        \"\"\"\n",
    "        model_device = next(self.parameters()).device\n",
    "        states = torch.tensor(states, device=model_device, dtype=torch.float32)\n",
    "        qvalues = self.forward(states)\n",
    "        return qvalues.data.cpu().numpy()\n",
    "\n",
    "    def sample_actions(self, qvalues):\n",
    "        \"\"\"pick actions given qvalues. Uses epsilon-greedy exploration strategy. \"\"\"\n",
    "        epsilon = self.epsilon\n",
    "        batch_size, n_actions = qvalues.shape\n",
    "\n",
    "        random_actions = np.random.choice(n_actions, size=batch_size)\n",
    "        best_actions = qvalues.argmax(axis=-1)\n",
    "\n",
    "        should_explore = np.random.choice(\n",
    "            [0, 1], batch_size, p=[1-epsilon, epsilon])\n",
    "        return np.where(should_explore, random_actions, best_actions)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "Bv1s5JKzZdC-"
   },
   "outputs": [],
   "source": [
    "agent = DQNAgent(state_shape, n_actions, epsilon=0.5).to(device)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "vazC0DPQZdC_"
   },
   "source": [
    "Now let's try out our agent to see if it raises any errors."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "e-Sg1cqPZdC_"
   },
   "outputs": [],
   "source": [
    "def evaluate(env, agent, n_games=1, greedy=False, t_max=10000, seed=None):\n",
    "    \"\"\" Plays n_games full games. If greedy, picks actions as argmax(qvalues). Returns mean reward. \"\"\"\n",
    "    rewards = []\n",
    "    for _ in range(n_games):\n",
    "        s, _ = env.reset(seed=seed)\n",
    "        reward = 0\n",
    "        for _ in range(t_max):\n",
    "            qvalues = agent.get_qvalues([s])\n",
    "            action = qvalues.argmax(axis=-1)[0] if greedy else agent.sample_actions(qvalues)[0]\n",
    "            s, r, terminated, truncated, _ = env.step(action)\n",
    "            reward += r\n",
    "            if terminated or truncated:\n",
    "                break\n",
    "\n",
    "        rewards.append(reward)\n",
    "    return np.mean(rewards)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Y_0NzjUEZdC_"
   },
   "source": [
    "### Experience replay\n",
    "For this assignment, we provide you with experience replay buffer. If you implemented experience replay buffer in last week's assignment, you can copy-paste it here in main notebook **to get 2 bonus points**.\n",
    "\n",
    "![img](https://github.com/yandexdataschool/Practical_RL/raw/master/yet_another_week/_resource/exp_replay.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "jHyCO4TuZdC_"
   },
   "source": [
    "#### The interface is fairly simple:\n",
    "* `exp_replay.add(obs, act, rw, next_obs, done)` - saves (s,a,r,s',done) tuple into the buffer\n",
    "* `exp_replay.sample(batch_size)` - returns observations, actions, rewards, next_observations and is_done for `batch_size` random samples.\n",
    "* `len(exp_replay)` - returns number of elements stored in replay buffer."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "wQEHwR1AZdC_"
   },
   "outputs": [],
   "source": [
    "from dqn.replay_buffer import ReplayBuffer\n",
    "exp_replay = ReplayBuffer(10)\n",
    "\n",
    "for _ in range(30):\n",
    "    exp_replay.add(env.reset()[0], env.action_space.sample(), 1.0, env.reset()[0], done=False)\n",
    "\n",
    "obs_batch, act_batch, reward_batch, next_obs_batch, is_done_batch = exp_replay.sample(5)\n",
    "\n",
    "assert len(exp_replay) == 10, \"experience replay size should be 10 because that's what maximum capacity is\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "0RnFX5sfZdC_"
   },
   "outputs": [],
   "source": [
    "def play_and_record(initial_state, agent, env, exp_replay, n_steps=1):\n",
    "    \"\"\"\n",
    "    Play the game for exactly n_steps, record every (s,a,r,s', done) to replay buffer.\n",
    "    Whenever game ends due to termination or truncation, add record with done=terminated and reset the game.\n",
    "    It is guaranteed that env has terminated=False when passed to this function.\n",
    "\n",
    "    PLEASE DO NOT RESET ENV UNLESS IT IS \"DONE\"\n",
    "\n",
    "    :returns: return sum of rewards over time and the state in which the env stays\n",
    "    \"\"\"\n",
    "    s = initial_state\n",
    "    sum_rewards = 0\n",
    "\n",
    "    # Play the game for n_steps as per instructions above\n",
    "    <YOUR CODE>\n",
    "\n",
    "    return sum_rewards, s"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "ZXXmFEKGZdC_",
    "outputId": "d1b66847-a141-4406-9697-7ebd194fdb6a"
   },
   "outputs": [],
   "source": [
    "# testing your code.\n",
    "exp_replay = ReplayBuffer(2000)\n",
    "\n",
    "state, _ = env.reset()\n",
    "play_and_record(state, agent, env, exp_replay, n_steps=1000)\n",
    "\n",
    "# if you're using your own experience replay buffer, some of those tests may need correction.\n",
    "# just make sure you know what your code does\n",
    "assert len(exp_replay) == 1000, \\\n",
    "    \"play_and_record should have added exactly 1000 steps, \" \\\n",
    "    \"but instead added %i\" % len(exp_replay)\n",
    "is_dones = list(zip(*exp_replay._storage))[-1]\n",
    "\n",
    "assert 0 < np.mean(is_dones) < 0.1, \\\n",
    "    \"Please make sure you restart the game whenever it is 'done' and \" \\\n",
    "    \"record the is_done correctly into the buffer. Got %f is_done rate over \" \\\n",
    "    \"%i steps. [If you think it's your tough luck, just re-run the test]\" % (\n",
    "        np.mean(is_dones), len(exp_replay))\n",
    "\n",
    "for _ in range(100):\n",
    "    obs_batch, act_batch, reward_batch, next_obs_batch, is_done_batch = exp_replay.sample(10)\n",
    "    assert obs_batch.shape == next_obs_batch.shape == (10,) + state_shape\n",
    "    assert act_batch.shape == (10,), \\\n",
    "        \"actions batch should have shape (10,) but is instead %s\" % str(act_batch.shape)\n",
    "    assert reward_batch.shape == (10,), \\\n",
    "        \"rewards batch should have shape (10,) but is instead %s\" % str(reward_batch.shape)\n",
    "    assert is_done_batch.shape == (10,), \\\n",
    "        \"is_done batch should have shape (10,) but is instead %s\" % str(is_done_batch.shape)\n",
    "    assert [int(i) in (0, 1) for i in is_dones], \\\n",
    "        \"is_done should be strictly True or False\"\n",
    "    assert [0 <= a < n_actions for a in act_batch], \"actions should be within [0, n_actions)\"\n",
    "\n",
    "print(\"Well done!\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "uoVGsnHRZdC_"
   },
   "source": [
    "### Target networks\n",
    "\n",
    "We also employ the so called \"target network\" - a copy of neural network weights to be used for reference Q-values:\n",
    "\n",
    "The network itself is an exact copy of agent network, but it's parameters are not trained. Instead, they are moved here from agent's actual network every so often.\n",
    "\n",
    "$$ Q_{reference}(s,a) = r + \\gamma \\cdot \\max _{a'} Q_{target}(s',a') $$\n",
    "\n",
    "![img](https://github.com/yandexdataschool/Practical_RL/raw/master/yet_another_week/_resource/target_net.png)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "8BLJCNiuZdC_",
    "outputId": "6181261a-60cf-4626-fbe6-930a6ccd9896"
   },
   "outputs": [],
   "source": [
    "target_network = DQNAgent(agent.state_shape, agent.n_actions, epsilon=0.5).to(device)\n",
    "# This is how you can load weights from agent into target network\n",
    "target_network.load_state_dict(agent.state_dict())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "I_GGShX3ZdC_"
   },
   "source": [
    "### Learning with... Q-learning\n",
    "Here we write a function similar to `agent.update` from tabular q-learning."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "4hbg-xANZdC_"
   },
   "source": [
    "Compute Q-learning TD error:\n",
    "\n",
    "$$ L = { 1 \\over N} \\sum_i [ Q_{\\theta}(s,a) - Q_{reference}(s,a) ] ^2 $$\n",
    "\n",
    "With Q-reference defined as\n",
    "\n",
    "$$ Q_{reference}(s,a) = r(s,a) + \\gamma \\cdot max_{a'} Q_{target}(s', a') $$\n",
    "\n",
    "Where\n",
    "* $Q_{target}(s',a')$ denotes Q-value of next state and next action predicted by __target_network__\n",
    "* $s, a, r, s'$ are current state, action, reward and next state respectively\n",
    "* $\\gamma$ is a discount factor defined two cells above.\n",
    "\n",
    "\n",
    "__Note 1:__ there's an example input below. Feel free to experiment with it before you write the function.\n",
    "\n",
    "__Note 2:__ compute_td_loss is a source of 99% of bugs in this homework. If reward doesn't improve, it often helps to go through it line by line [with a rubber duck](https://rubberduckdebugging.com/)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "VxrEOC7mZdC_"
   },
   "outputs": [],
   "source": [
    "def compute_td_loss(states, actions, rewards, next_states, is_done,\n",
    "                    agent, target_network,\n",
    "                    gamma=0.99,\n",
    "                    check_shapes=False,\n",
    "                    device=device):\n",
    "    \"\"\" Compute td loss using torch operations only. Use the formulae above. \"\"\"\n",
    "    states = torch.tensor(states, device=device, dtype=torch.float32)    # shape: [batch_size, *state_shape]\n",
    "    actions = torch.tensor(actions, device=device, dtype=torch.int64)    # shape: [batch_size]\n",
    "    rewards = torch.tensor(rewards, device=device, dtype=torch.float32)  # shape: [batch_size]\n",
    "    # shape: [batch_size, *state_shape]\n",
    "    next_states = torch.tensor(next_states, device=device, dtype=torch.float)\n",
    "    is_done = torch.tensor(\n",
    "        is_done.astype('float32'),\n",
    "        device=device,\n",
    "        dtype=torch.float32,\n",
    "    )  # shape: [batch_size]\n",
    "    is_not_done = 1 - is_done\n",
    "\n",
    "    # get q-values for all actions in current states\n",
    "    predicted_qvalues = agent(states)  # shape: [batch_size, n_actions]\n",
    "\n",
    "    # compute q-values for all actions in next states\n",
    "    predicted_next_qvalues = target_network(next_states)  # shape: [batch_size, n_actions]\n",
    "\n",
    "    # select q-values for chosen actions\n",
    "    predicted_qvalues_for_actions = predicted_qvalues[range(len(actions)), actions]  # shape: [batch_size]\n",
    "\n",
    "    # compute V*(next_states) using predicted next q-values\n",
    "    next_state_values = <YOUR CODE>\n",
    "\n",
    "    assert next_state_values.dim() == 1 and next_state_values.shape[0] == states.shape[0], \\\n",
    "        \"must predict one value per state\"\n",
    "\n",
    "    # compute \"target q-values\" for loss - it's what's inside square parentheses in the above formula.\n",
    "    # at the last state use the simplified formula: Q(s,a) = r(s,a) since s' doesn't exist\n",
    "    # you can multiply next state values by is_not_done to achieve this.\n",
    "    target_qvalues_for_actions = <YOUR CODE>\n",
    "\n",
    "    # mean squared error loss to minimize\n",
    "    loss = torch.mean((predicted_qvalues_for_actions - target_qvalues_for_actions.detach()) ** 2)\n",
    "\n",
    "    if check_shapes:\n",
    "        assert predicted_next_qvalues.data.dim() == 2, \\\n",
    "            \"make sure you predicted q-values for all actions in next state\"\n",
    "        assert next_state_values.data.dim() == 1, \\\n",
    "            \"make sure you computed V(s') as maximum over just the actions axis and not all axes\"\n",
    "        assert target_qvalues_for_actions.data.dim() == 1, \\\n",
    "            \"there's something wrong with target q-values, they must be a vector\"\n",
    "\n",
    "    return loss"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "pgZKcPPnZdC_"
   },
   "source": [
    "Sanity checks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "Yp8eREoDZdC_"
   },
   "outputs": [],
   "source": [
    "obs_batch, act_batch, reward_batch, next_obs_batch, is_done_batch = exp_replay.sample(10)\n",
    "\n",
    "loss = compute_td_loss(obs_batch, act_batch, reward_batch, next_obs_batch, is_done_batch,\n",
    "                       agent, target_network,\n",
    "                       gamma=0.99, check_shapes=True)\n",
    "loss.backward()\n",
    "\n",
    "assert loss.requires_grad and tuple(loss.data.size()) == (), \\\n",
    "    \"you must return scalar loss - mean over batch\"\n",
    "assert np.any(next(agent.parameters()).grad.data.cpu().numpy() != 0), \\\n",
    "    \"loss must be differentiable w.r.t. network weights\"\n",
    "assert np.all(next(target_network.parameters()).grad is None), \\\n",
    "    \"target network should not have grads\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "8A1QtGVqZdC_"
   },
   "source": [
    "### Main loop\n",
    "\n",
    "It's time to put everything together and see if it learns anything."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "8lAUT94JZdC_"
   },
   "outputs": [],
   "source": [
    "from tqdm import trange\n",
    "from IPython.display import clear_output\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "YOk81bdZZdC_",
    "outputId": "2fd2404e-19e5-4ebe-b0e1-c6f7593db790"
   },
   "outputs": [],
   "source": [
    "seed = <YOUR CODE: your favourite random seed>\n",
    "random.seed(seed)\n",
    "np.random.seed(seed)\n",
    "torch.manual_seed(seed)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "13K5t2CTZdDA",
    "outputId": "031d4a0a-99c3-4cc3-f7c3-77e4d8a5a331"
   },
   "outputs": [],
   "source": [
    "state_dim = env.observation_space.shape\n",
    "n_actions = env.action_space.n\n",
    "state, _ = env.reset(seed=seed)\n",
    "\n",
    "agent = DQNAgent(state_dim, n_actions, epsilon=1).to(device)\n",
    "target_network = DQNAgent(state_dim, n_actions, epsilon=1).to(device)\n",
    "target_network.load_state_dict(agent.state_dict())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "iD7PAlwQZdDA",
    "outputId": "aeb4bb67-4776-4b02-e558-d9d1a3306d47"
   },
   "outputs": [],
   "source": [
    "from dqn.utils import is_enough_ram, linear_decay\n",
    "\n",
    "REPLAY_BUFFER_SIZE = 10**4\n",
    "\n",
    "exp_replay = ReplayBuffer(REPLAY_BUFFER_SIZE)\n",
    "for i in range(100):\n",
    "    if not is_enough_ram(min_available_gb=0.1):\n",
    "        print(\"\"\"\n",
    "            Less than 100 Mb RAM available.\n",
    "            Make sure the buffer size in not too huge.\n",
    "            Also check, maybe other processes consume RAM heavily.\n",
    "            \"\"\"\n",
    "             )\n",
    "        break\n",
    "    play_and_record(state, agent, env, exp_replay, n_steps=10**2)\n",
    "    if len(exp_replay) == REPLAY_BUFFER_SIZE:\n",
    "        break\n",
    "print(len(exp_replay))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "Zl2VCEYQZdDA"
   },
   "outputs": [],
   "source": [
    "# # for something more complicated than CartPole\n",
    "\n",
    "# timesteps_per_epoch = 1\n",
    "# batch_size = 32\n",
    "# total_steps = 3 * 10**6\n",
    "# decay_steps = 1 * 10**6\n",
    "\n",
    "# opt = torch.optim.Adam(agent.parameters(), lr=1e-4)\n",
    "\n",
    "# init_epsilon = 1\n",
    "# final_epsilon = 0.1\n",
    "\n",
    "# loss_freq = 20\n",
    "# refresh_target_network_freq = 1000\n",
    "# eval_freq = 5000\n",
    "\n",
    "# max_grad_norm = 5000"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "x-sD-QyUZdDA"
   },
   "outputs": [],
   "source": [
    "timesteps_per_epoch = 1\n",
    "batch_size = 32\n",
    "total_steps = 4 * 10**4\n",
    "decay_steps = 1 * 10**4\n",
    "\n",
    "opt = torch.optim.Adam(agent.parameters(), lr=1e-4)\n",
    "\n",
    "init_epsilon = 1\n",
    "final_epsilon = 0.1\n",
    "\n",
    "loss_freq = 20\n",
    "refresh_target_network_freq = 100\n",
    "eval_freq = 1000\n",
    "\n",
    "max_grad_norm = 5000"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "piqDfKQAZdDA"
   },
   "outputs": [],
   "source": [
    "mean_rw_history = []\n",
    "td_loss_history = []\n",
    "grad_norm_history = []\n",
    "initial_state_v_history = []\n",
    "step = 0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "Ks8NAV8AZdDA"
   },
   "outputs": [],
   "source": [
    "import time\n",
    "\n",
    "def wait_for_keyboard_interrupt():\n",
    "    try:\n",
    "        while True:\n",
    "            time.sleep(1)\n",
    "    except KeyboardInterrupt:\n",
    "        pass"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "sU3GSGZqZdDA"
   },
   "outputs": [],
   "source": [
    "state, _ = env.reset()\n",
    "with trange(step, total_steps + 1) as progress_bar:\n",
    "    for step in progress_bar:\n",
    "        if not is_enough_ram():\n",
    "            print('less that 100 Mb RAM available, freezing')\n",
    "            print('make sure everything is ok and use KeyboardInterrupt to continue')\n",
    "            wait_for_keyboard_interrupt()\n",
    "\n",
    "        agent.epsilon = linear_decay(init_epsilon, final_epsilon, step, decay_steps)\n",
    "\n",
    "        # play\n",
    "        _, state = play_and_record(state, agent, env, exp_replay, timesteps_per_epoch)\n",
    "\n",
    "        # train\n",
    "        <YOUR CODE: sample batch_size of data from experience replay>\n",
    "\n",
    "        loss = <YOUR CODE: compute TD loss>\n",
    "\n",
    "        loss.backward()\n",
    "        grad_norm = nn.utils.clip_grad_norm_(agent.parameters(), max_grad_norm)\n",
    "        opt.step()\n",
    "        opt.zero_grad()\n",
    "\n",
    "        if step % loss_freq == 0:\n",
    "            td_loss_history.append(loss.data.cpu().item())\n",
    "            grad_norm_history.append(grad_norm)\n",
    "\n",
    "        if step % refresh_target_network_freq == 0:\n",
    "            # Load agent weights into target_network\n",
    "            <YOUR CODE>\n",
    "\n",
    "        if step % eval_freq == 0:\n",
    "            mean_rw_history.append(evaluate(\n",
    "                make_env(), agent, n_games=3, greedy=True, t_max=1000, seed=step)\n",
    "            )\n",
    "            initial_state_q_values = agent.get_qvalues(\n",
    "                [make_env().reset(seed=step)[0]]\n",
    "            )\n",
    "            initial_state_v_history.append(np.max(initial_state_q_values))\n",
    "\n",
    "            clear_output(True)\n",
    "            print(\"buffer size = %i, epsilon = %.5f\" %\n",
    "                (len(exp_replay), agent.epsilon))\n",
    "\n",
    "            plt.figure(figsize=[16, 9])\n",
    "\n",
    "            plt.subplot(2, 2, 1)\n",
    "            plt.title(\"Mean reward per episode\")\n",
    "            plt.plot(mean_rw_history)\n",
    "            plt.grid()\n",
    "\n",
    "            assert not np.isnan(td_loss_history[-1])\n",
    "            plt.subplot(2, 2, 2)\n",
    "            plt.title(\"TD loss history\")\n",
    "            plt.plot(td_loss_history)\n",
    "            plt.grid()\n",
    "\n",
    "            plt.subplot(2, 2, 3)\n",
    "            plt.title(\"Initial state V\")\n",
    "            plt.plot(initial_state_v_history)\n",
    "            plt.grid()\n",
    "\n",
    "            plt.subplot(2, 2, 4)\n",
    "            plt.title(\"Grad norm history\")\n",
    "            plt.plot(grad_norm_history)\n",
    "            plt.grid()\n",
    "\n",
    "            plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "qwWFT2SBZdDA"
   },
   "outputs": [],
   "source": [
    "final_score = evaluate(\n",
    "  make_env(),\n",
    "  agent, n_games=30, greedy=True, t_max=1000\n",
    ")\n",
    "print('final score:', final_score)\n",
    "assert final_score > 300, 'not good enough for DQN'\n",
    "print('Well done')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "G-feeX9YZdDA"
   },
   "source": [
    "**Agent's predicted V-values vs their Monte-Carlo estimates**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "rjVuSIrPZdDA"
   },
   "outputs": [],
   "source": [
    "from analysis import play_and_log_episode\n",
    "\n",
    "eval_env = make_env()\n",
    "record = play_and_log_episode(eval_env, agent)\n",
    "print('total reward for life:', np.sum(record['rewards']))\n",
    "for key in record:\n",
    "    print(key)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "FCacwLw6ZdDA"
   },
   "outputs": [],
   "source": [
    "fig = plt.figure(figsize=(5, 5))\n",
    "ax = fig.add_subplot(1, 1, 1)\n",
    "\n",
    "ax.scatter(record['v_mc'], record['v_agent'])\n",
    "ax.plot(sorted(record['v_mc']), sorted(record['v_mc']),\n",
    "       'black', linestyle='--', label='x=y')\n",
    "\n",
    "ax.grid()\n",
    "ax.legend()\n",
    "ax.set_title('State Value Estimates')\n",
    "ax.set_xlabel('Monte-Carlo')\n",
    "ax.set_ylabel('Agent')\n",
    "\n",
    "plt.show()"
   ]
  }
 ],
 "metadata": {
  "colab": {
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3",
   "name": "python3"
  },
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}