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",
        "    !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\n",
        "    !wget -q https://raw.githubusercontent.com/yandexdataschool/Practical_RL/master/week04_approx_rl/dqn/utils.py\n",
        "    !wget -q https://raw.githubusercontent.com/yandexdataschool/Practical_RL/master/week04_approx_rl/dqn/replay_buffer.py\n",
        "    !wget -q https://raw.githubusercontent.com/yandexdataschool/Practical_RL/master/week04_approx_rl/dqn/framebuffer.py\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": "code",
      "execution_count": null,
      "metadata": {
        "id": "u8OFQOtGojc8"
      },
      "outputs": [],
      "source": [
        "!pip install gymnasium"
      ]
    },
    {
      "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 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 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": [
        "REPLAY_BUFFER_SIZE = 10**4\n",
        "\n",
        "exp_replay = ReplayBuffer(REPLAY_BUFFER_SIZE)\n",
        "for i in range(100):\n",
        "    if not utils.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 utils.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 = utils.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 (smoothened)\")\n",
        "            plt.plot(utils.smoothen(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 (smoothened)\")\n",
        "            plt.plot(utils.smoothen(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": [
        "eval_env = make_env()\n",
        "record = utils.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
}