FlatLand_DQN.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "FlatLand DQN.ipynb",
      "provenance": [],
      "collapsed_sections": [],
      "authorship_tag": "ABX9TyN3wDyCyN9GZq2LrQBtH8uH",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/FatLads/Notebooks/blob/main/FlatLand_DQN.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "ZfdEAgKjB2ay",
        "outputId": "e436f433-4cce-4579-f4b6-4d91b3152436"
      },
      "source": [
        "!sudo apt install -y xvfb ffmpeg\n",
        "!pip install -q 'imageio==2.4.0'\n",
        "!pip install -q pyvirtualdisplay\n",
        "!pip install -q tf-agents\n",
        "!pip install -q flatland-rl"
      ],
      "execution_count": 1,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Reading package lists... Done\n",
            "Building dependency tree       \n",
            "Reading state information... Done\n",
            "ffmpeg is already the newest version (7:3.4.8-0ubuntu0.2).\n",
            "The following NEW packages will be installed:\n",
            "  xvfb\n",
            "0 upgraded, 1 newly installed, 0 to remove and 29 not upgraded.\n",
            "Need to get 784 kB of archives.\n",
            "After this operation, 2,270 kB of additional disk space will be used.\n",
            "Get:1 http://archive.ubuntu.com/ubuntu bionic-updates/universe amd64 xvfb amd64 2:1.19.6-1ubuntu4.8 [784 kB]\n",
            "Fetched 784 kB in 1s (703 kB/s)\n",
            "debconf: unable to initialize frontend: Dialog\n",
            "debconf: (No usable dialog-like program is installed, so the dialog based frontend cannot be used. at /usr/share/perl5/Debconf/FrontEnd/Dialog.pm line 76, <> line 1.)\n",
            "debconf: falling back to frontend: Readline\n",
            "debconf: unable to initialize frontend: Readline\n",
            "debconf: (This frontend requires a controlling tty.)\n",
            "debconf: falling back to frontend: Teletype\n",
            "dpkg-preconfigure: unable to re-open stdin: \n",
            "Selecting previously unselected package xvfb.\n",
            "(Reading database ... 160975 files and directories currently installed.)\n",
            "Preparing to unpack .../xvfb_2%3a1.19.6-1ubuntu4.8_amd64.deb ...\n",
            "Unpacking xvfb (2:1.19.6-1ubuntu4.8) ...\n",
            "Setting up xvfb (2:1.19.6-1ubuntu4.8) ...\n",
            "Processing triggers for man-db (2.8.3-2ubuntu0.1) ...\n",
            "\u001b[K     |████████████████████████████████| 3.3MB 7.8MB/s \n",
            "\u001b[?25h  Building wheel for imageio (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
            "\u001b[31mERROR: albumentations 0.1.12 has requirement imgaug<0.2.7,>=0.2.5, but you'll have imgaug 0.2.9 which is incompatible.\u001b[0m\n",
            "\u001b[K     |████████████████████████████████| 1.2MB 9.0MB/s \n",
            "\u001b[K     |████████████████████████████████| 3.3MB 8.7MB/s \n",
            "\u001b[?25h  Building wheel for flatland-rl (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
            "  Building wheel for crowdai-api (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
            "  Building wheel for timeout-decorator (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
            "  Building wheel for gym (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
            "\u001b[31mERROR: tf-agents 0.7.1 has requirement cloudpickle>=1.3, but you'll have cloudpickle 1.2.2 which is incompatible.\u001b[0m\n",
            "\u001b[31mERROR: tf-agents 0.7.1 has requirement gym>=0.17.0, but you'll have gym 0.14.0 which is incompatible.\u001b[0m\n",
            "\u001b[31mERROR: tensorflow-probability 0.12.1 has requirement cloudpickle>=1.3, but you'll have cloudpickle 1.2.2 which is incompatible.\u001b[0m\n",
            "\u001b[31mERROR: datascience 0.10.6 has requirement folium==0.2.1, but you'll have folium 0.8.3 which is incompatible.\u001b[0m\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "CxiXt5sTB9PK"
      },
      "source": [
        "from __future__ import absolute_import, division, print_function\n",
        "\n",
        "import base64\n",
        "import imageio\n",
        "import IPython\n",
        "import matplotlib\n",
        "import matplotlib.pyplot as plt\n",
        "import numpy as np\n",
        "import PIL.Image\n",
        "import pyvirtualdisplay\n",
        "\n",
        "import tensorflow as tf\n",
        "\n",
        "from tf_agents.agents.dqn import dqn_agent\n",
        "from tf_agents.drivers import dynamic_step_driver\n",
        "from tf_agents.environments import suite_gym\n",
        "from tf_agents.environments import tf_py_environment, py_environment\n",
        "from tf_agents.eval import metric_utils\n",
        "from tf_agents.networks import sequential \n",
        "from tf_agents.policies import random_tf_policy \n",
        "from tf_agents.environments import utils\n",
        "from tf_agents.replay_buffers import tf_uniform_replay_buffer\n",
        "from tf_agents.trajectories import trajectory, time_step\n",
        "from tf_agents.specs import tensor_spec, BoundedArraySpec\n",
        "from tf_agents.utils import common \n",
        "from flatland.envs.rail_env import RailEnv, RailEnvActions\n",
        "\n",
        "display = pyvirtualdisplay.Display(visible=0, size=(1400,900)).start()"
      ],
      "execution_count": 9,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "_ZVaFvxQCBP-"
      },
      "source": [
        "num_iterations = 30000 # @param {type:\"integer\"}\n",
        "\n",
        "initial_collect_steps = 100 # @param {type:\"integer\"}\n",
        "collect_steps_per_iteration = 1 # @param {type:\"integer\"}\n",
        "replay_buffer_max_length = 100000 # @param {type:\"integer\"}\n",
        "\n",
        "batch_size = 64 # @param {type:\"integer\"}\n",
        "learning_rate = 1e-4 # @param {type:\"number\"}\n",
        "log_interval = 200 # @param {type:\"integer\"}\n",
        "\n",
        "num_eval_episodes = 10 # @param {type:\"integer\"}\n",
        "eval_interval = 1000 # @param {type:\"integer\"}"
      ],
      "execution_count": 3,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "J3l9WfEDh_Gu"
      },
      "source": [
        "Next, we're gonna try and reimplement the `Rails` env with all of the things that TensorFlow needs, like explained [here](https://www.tensorflow.org/agents/tutorials/2_environments_tutorial) "
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "cwmLo4iah9z0"
      },
      "source": [
        "class FatRails(py_environment.PyEnvironment):\n",
        "    def __init__(self, *args, **kwargs):\n",
        "        self.env = RailEnv(*args, **kwargs)\n",
        "        self._episode_ended = False\n",
        "\n",
        "    def action_spec(self):\n",
        "        return BoundedArraySpec(shape=(), dtype=np.int32, minimum=0, maximum=4, name='action') # We need to convert from an array of actions (index=agent, value=action) to a dict\n",
        "    def observation_spec(self):\n",
        "        return BoundedArraySpec(shape=(self.env.width, self.env.height, 16), dtype=np.int32, minimum=0, maximum=1, name='observation') #TODO\n",
        "    def _step(self, action):\n",
        "        action_dict = {v: k for v, k in enumerate(action)}\n",
        "        step_env = self.env.step(action_dict)\n",
        "        agents_statuses = step_env[3]['status']\n",
        "        for status in agents_statuses.values(): # Check if there's someone that didn't arrive yet\n",
        "            if status<3: # The status observations are 0..3, check the docs\n",
        "                return time_step.transition(step_env[0][0][0], step_env[1][0])\n",
        "        self._episode_ended = True\n",
        "        return time_step.termination(step_env[0][0][0], step_env[1][0]) # If no one is moving/has to depart, we're finished\n",
        "    def _reset(self):\n",
        "        reset = self.env.reset()\n",
        "        self._episode_ended = False\n",
        "        return time_step.restart(np.array(reset[0][0][0], dtype=np.int32))"
      ],
      "execution_count": 34,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "r8A__-LZ95Bk"
      },
      "source": [
        "Note the `status<3`: this exploits the agents' statuses we can find [here](https://gitlab.aicrowd.com/flatland/flatland/-/blob/master/flatland/envs/agent_utils.py)"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "cip9mv0Tx3D_"
      },
      "source": [
        "env = FatRails(16,16)"
      ],
      "execution_count": 35,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "7P-lEpk_4-df",
        "outputId": "cdfeee92-4f1a-4541-e894-90eb5dfcb166"
      },
      "source": [
        "env.reset()\n",
        "env.step([RailEnvActions.MOVE_FORWARD])"
      ],
      "execution_count": 36,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "TimeStep(step_type=array(1, dtype=int32), reward=array(-1., dtype=float32), discount=array(1., dtype=float32), observation=array([[[0., 0., 0., ..., 0., 0., 0.],\n",
              "        [0., 0., 1., ..., 0., 0., 0.],\n",
              "        [0., 0., 0., ..., 0., 0., 0.],\n",
              "        ...,\n",
              "        [0., 0., 1., ..., 0., 0., 0.],\n",
              "        [0., 0., 1., ..., 0., 0., 0.],\n",
              "        [0., 0., 0., ..., 0., 0., 0.]],\n",
              "\n",
              "       [[0., 0., 0., ..., 1., 0., 0.],\n",
              "        [0., 0., 0., ..., 0., 0., 0.],\n",
              "        [0., 1., 0., ..., 0., 1., 0.],\n",
              "        ...,\n",
              "        [0., 0., 0., ..., 0., 0., 1.],\n",
              "        [1., 0., 0., ..., 0., 0., 0.],\n",
              "        [0., 0., 0., ..., 0., 0., 0.]],\n",
              "\n",
              "       [[0., 0., 0., ..., 1., 0., 0.],\n",
              "        [1., 0., 0., ..., 0., 0., 1.],\n",
              "        [0., 0., 0., ..., 0., 0., 0.],\n",
              "        ...,\n",
              "        [0., 1., 0., ..., 0., 1., 0.],\n",
              "        [1., 0., 0., ..., 0., 0., 1.],\n",
              "        [0., 0., 0., ..., 0., 0., 0.]],\n",
              "\n",
              "       ...,\n",
              "\n",
              "       [[0., 0., 0., ..., 1., 0., 0.],\n",
              "        [1., 0., 0., ..., 0., 0., 0.],\n",
              "        [0., 0., 0., ..., 0., 0., 0.],\n",
              "        ...,\n",
              "        [1., 1., 0., ..., 0., 1., 0.],\n",
              "        [1., 0., 0., ..., 0., 0., 0.],\n",
              "        [0., 0., 0., ..., 0., 0., 0.]],\n",
              "\n",
              "       [[0., 0., 0., ..., 1., 0., 0.],\n",
              "        [1., 1., 0., ..., 0., 1., 1.],\n",
              "        [0., 0., 0., ..., 0., 0., 0.],\n",
              "        ...,\n",
              "        [0., 0., 0., ..., 0., 0., 0.],\n",
              "        [1., 0., 0., ..., 0., 0., 1.],\n",
              "        [0., 0., 0., ..., 0., 0., 0.]],\n",
              "\n",
              "       [[0., 0., 0., ..., 0., 0., 0.],\n",
              "        [0., 0., 0., ..., 0., 0., 0.],\n",
              "        [0., 0., 0., ..., 0., 0., 0.],\n",
              "        ...,\n",
              "        [0., 0., 0., ..., 0., 0., 0.],\n",
              "        [0., 0., 0., ..., 0., 0., 0.],\n",
              "        [0., 0., 0., ..., 0., 0., 0.]]]))"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 36
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "oPhBk25CrkZJ"
      },
      "source": [
        "train_env = tf_py_environment.TFPyEnvironment(env)\n",
        "test_env = tf_py_environment.TFPyEnvironment(env)"
      ],
      "execution_count": 38,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "me-jnQnc_39Q",
        "outputId": "927b1430-790d-4484-9dda-7e6427b2d280"
      },
      "source": [
        "train_env.time_step_spec()"
      ],
      "execution_count": 39,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "TimeStep(step_type=TensorSpec(shape=(), dtype=tf.int32, name='step_type'), reward=TensorSpec(shape=(), dtype=tf.float32, name='reward'), discount=BoundedTensorSpec(shape=(), dtype=tf.float32, name='discount', minimum=array(0., dtype=float32), maximum=array(1., dtype=float32)), observation=BoundedTensorSpec(shape=(16, 16, 16), dtype=tf.int32, name='observation', minimum=array(0, dtype=int32), maximum=array(1, dtype=int32)))"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 39
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "OU_NryIcAV4k",
        "outputId": "f90c7718-dc6e-4476-a5c3-2f94ba15771b"
      },
      "source": [
        "action_tensor_spec = tensor_spec.from_spec(env.action_spec())\n",
        "action_tensor_spec.maximum-action_tensor_spec.minimum + 1\n"
      ],
      "execution_count": 40,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "5"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 40
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "EehoB0wZDcJk"
      },
      "source": [
        "fc_layer_params = (100,50)\n",
        "\n",
        "action_tensor_spec = tensor_spec.from_spec(env.action_spec())\n",
        "num_actions = action_tensor_spec.maximum-action_tensor_spec.minimum + 1\n",
        "\n",
        "\n",
        "\n",
        "def dense_layer(num_units):\n",
        "    return tf.keras.layers.Dense(\n",
        "        num_units,\n",
        "        activation=tf.keras.activations.relu,\n",
        "        kernel_initializer=tf.keras.initializers.VarianceScaling(\n",
        "            scale=2.0, mode='fan_in', distribution='truncated_normal'\n",
        "        )\n",
        "    )\n",
        "\n",
        "dense_layers = [dense_layer(num_units) for num_units in fc_layer_params]\n",
        "q_values_layer = tf.keras.layers.Dense(\n",
        "    num_actions,\n",
        "    activation=None,\n",
        "    kernel_initializer=tf.keras.initializers.RandomUniform(\n",
        "        minval=-0.03, maxval=0.03\n",
        "    ),\n",
        "    bias_initializer=tf.keras.initializers.Constant(-0.2)\n",
        ")\n",
        "q_net = sequential.Sequential(dense_layers+[q_values_layer])"
      ],
      "execution_count": 46,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "H5JGOFUPAuhH"
      },
      "source": [
        ""
      ],
      "execution_count": 41,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 375
        },
        "id": "e_jha6EW8b9p",
        "outputId": "b8ca96c2-5542-4604-f865-3b4b0bc60bd6"
      },
      "source": [
        "optimizer = tf.keras.optimizers.Adam(learning_rate=learning_rate)\n",
        "train_step_counter = tf.Variable(0)\n",
        "agent = dqn_agent.DqnAgent(train_env.time_step_spec(),\n",
        "                           train_env.action_spec(),\n",
        "                           q_network=q_net,\n",
        "                           optimizer=optimizer,\n",
        "                           td_errors_loss_fn=common.element_wise_squared_loss,\n",
        "                           train_step_counter=train_step_counter)\n",
        "agent.initialize()"
      ],
      "execution_count": 47,
      "outputs": [
        {
          "output_type": "error",
          "ename": "ValueError",
          "evalue": "ignored",
          "traceback": [
            "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
            "\u001b[0;31mValueError\u001b[0m                                Traceback (most recent call last)",
            "\u001b[0;32m<ipython-input-47-b6c2fa09f7fd>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m()\u001b[0m\n\u001b[1;32m      6\u001b[0m                            \u001b[0moptimizer\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0moptimizer\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      7\u001b[0m                            \u001b[0mtd_errors_loss_fn\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mcommon\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0melement_wise_squared_loss\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m----> 8\u001b[0;31m                            train_step_counter=train_step_counter)\n\u001b[0m\u001b[1;32m      9\u001b[0m \u001b[0magent\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0minitialize\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.7/dist-packages/gin/config.py\u001b[0m in \u001b[0;36mgin_wrapper\u001b[0;34m(*args, **kwargs)\u001b[0m\n\u001b[1;32m   1067\u001b[0m       \u001b[0mscope_info\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m\" in scope '{}'\"\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mformat\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mscope_str\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;32mif\u001b[0m \u001b[0mscope_str\u001b[0m \u001b[0;32melse\u001b[0m \u001b[0;34m''\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1068\u001b[0m       \u001b[0merr_str\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0merr_str\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mformat\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mname\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mfn_or_cls\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mscope_info\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 1069\u001b[0;31m       \u001b[0mutils\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0maugment_exception_message_and_reraise\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0me\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0merr_str\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1070\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1071\u001b[0m   \u001b[0;32mreturn\u001b[0m \u001b[0mgin_wrapper\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.7/dist-packages/gin/utils.py\u001b[0m in \u001b[0;36maugment_exception_message_and_reraise\u001b[0;34m(exception, message)\u001b[0m\n\u001b[1;32m     39\u001b[0m   \u001b[0mproxy\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mExceptionProxy\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     40\u001b[0m   \u001b[0mExceptionProxy\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m__qualname__\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mtype\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mexception\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m__qualname__\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 41\u001b[0;31m   \u001b[0;32mraise\u001b[0m \u001b[0mproxy\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mwith_traceback\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mexception\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m__traceback__\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;32mfrom\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     42\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     43\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.7/dist-packages/gin/config.py\u001b[0m in \u001b[0;36mgin_wrapper\u001b[0;34m(*args, **kwargs)\u001b[0m\n\u001b[1;32m   1044\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1045\u001b[0m     \u001b[0;32mtry\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 1046\u001b[0;31m       \u001b[0;32mreturn\u001b[0m \u001b[0mfn\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0mnew_args\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mnew_kwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1047\u001b[0m     \u001b[0;32mexcept\u001b[0m \u001b[0mException\u001b[0m \u001b[0;32mas\u001b[0m \u001b[0me\u001b[0m\u001b[0;34m:\u001b[0m  \u001b[0;31m# pylint: disable=broad-except\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1048\u001b[0m       \u001b[0merr_str\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m''\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.7/dist-packages/tf_agents/agents/dqn/dqn_agent.py\u001b[0m in \u001b[0;36m__init__\u001b[0;34m(self, time_step_spec, action_spec, q_network, optimizer, observation_and_action_constraint_splitter, epsilon_greedy, n_step_update, boltzmann_temperature, emit_log_probability, target_q_network, target_update_tau, target_update_period, td_errors_loss_fn, gamma, reward_scale_factor, gradient_clipping, debug_summaries, summarize_grads_and_vars, train_step_counter, name)\u001b[0m\n\u001b[1;32m    238\u001b[0m         name='TargetQNetwork')\n\u001b[1;32m    239\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 240\u001b[0;31m     \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_check_network_output\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_q_network\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m'q_network'\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    241\u001b[0m     \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_check_network_output\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_target_q_network\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m'target_q_network'\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    242\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.7/dist-packages/tf_agents/agents/dqn/dqn_agent.py\u001b[0m in \u001b[0;36m_check_network_output\u001b[0;34m(self, net, label)\u001b[0m\n\u001b[1;32m    321\u001b[0m         \u001b[0mnet\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcreate_variables\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    322\u001b[0m         \u001b[0mexpected_output_shape\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_num_actions\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 323\u001b[0;31m         label=label)\n\u001b[0m\u001b[1;32m    324\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    325\u001b[0m   def _setup_policy(self, time_step_spec, action_spec,\n",
            "\u001b[0;32m/usr/local/lib/python3.7/dist-packages/tf_agents/networks/utils.py\u001b[0m in \u001b[0;36mcheck_single_floating_network_output\u001b[0;34m(output_spec, expected_output_shape, label)\u001b[0m\n\u001b[1;32m     39\u001b[0m         \u001b[0;34m'Expected {} to emit a floating point tensor with inner dims '\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     40\u001b[0m         \u001b[0;34m'{}; but saw network output spec: {}'\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 41\u001b[0;31m         .format(label, expected_output_shape, output_spec))\n\u001b[0m\u001b[1;32m     42\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     43\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;31mValueError\u001b[0m: Expected q_network to emit a floating point tensor with inner dims (5,); but saw network output spec: TensorSpec(shape=(16, 16, 5), dtype=tf.float32, name=None)\n  In call to configurable 'DqnAgent' (<class 'tf_agents.agents.dqn.dqn_agent.DqnAgent'>)"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Wd_kLJ_o-ZgK"
      },
      "source": [
        "eval_policy = agent.policy\n",
        "collect_policy = agent.collect_policy"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "B8q8Gyv--a-_"
      },
      "source": [
        "def compute_avg_return(environment, policy, num_episodes):\n",
        "    total_return = 0.0\n",
        "    for _ in range(num_episodes):\n",
        "        time_step = environment.reset()\n",
        "        episode_return = 0.0\n",
        "        while not time_step.is_last():\n",
        "            action_step = policy.action(time_step)\n",
        "            time_step = environment.step(action_step)\n",
        "            episode_return += time_step.reward\n",
        "        total_return += episode_return\n",
        "\n",
        "    avg_return = total_return / num_episodes \n",
        "    return avg_return.numpy()[0]"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "QzswUcRs-dI9"
      },
      "source": [
        "replay_buffer = tf_uniform_replay_buffer.TFUniformReplayBuffer(\n",
        "    data_spec=agent.collect_data_spec,\n",
        "    batch_size=train_env.batch_size,\n",
        "    max_length=replay_buffer_max_length\n",
        ")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "WFK-qqEJ-erL"
      },
      "source": [
        "def collect_step(environment, policy, replay):\n",
        "    time_step = environment.current_time_step()\n",
        "    action_step = policy.action(time_step)\n",
        "    next_time_step = environment.step(action_step.action)\n",
        "    traj = trajectory.from_transition(time_step, action_step, next_time_step)\n",
        "    replay.add_batch(traj)\n",
        "\n",
        "def collect_data(env, policy, replay, steps):\n",
        "    for _ in range(steps):\n",
        "        collect_step(env, policy, replay)\n",
        "\n",
        "collect_data(train_env, random_policy, replay_buffer, initial_collect_steps)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Ht2rcc-f-gXu"
      },
      "source": [
        "data = replay_buffer.as_dataset(\n",
        "    num_parallel_calls=3,\n",
        "    sample_batch_size=batch_size,\n",
        "    num_steps=2\n",
        ").prefetch(3)\n",
        "\n",
        "iterator = iter(data)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "YcsQg8Qt-j6i"
      },
      "source": [
        "try:\n",
        "    %%time\n",
        "except:\n",
        "    pass\n",
        "\n",
        "agent.train = common.function(agent.train)\n",
        "\n",
        "agent.train_step_counter.assign(0)\n",
        "avg_return = compute_avg_return(test_env, agent.policy, num_eval_episodes)\n",
        "returns = [avg_return]\n",
        "\n",
        "for _ in range(5000):\n",
        "    old_data = data\n",
        "    collect_data(train_env, agent.collect_policy, replay_buffer, collect_steps_per_iteration)\n",
        "    experience, unused_info = next(iterator)\n",
        "    train_loss = agent.train(experience).loss\n",
        "\n",
        "    step = agent.train_step_counter.numpy()\n",
        "\n",
        "    if step % log_interval == 0:\n",
        "        print(f\"step={step}: loss = {train_loss}\")\n",
        "    if step % eval_interval ==0:\n",
        "        avg_return = compute_avg_return(test_env, agent.policy, num_eval_episodes)\n",
        "        print(f\"Step {step}, avg_ret: {avg_return}\")\n",
        "        returns.append(avg_return)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "TAFQJK1b-mfY"
      },
      "source": [
        "iterations = range(0, num_iterations+1, eval_interval)\n",
        "plt.plot(iterations, returns)\n",
        "plt.ylabel('Average Return')\n",
        "plt.xlabel('Iterations')\n",
        "plt.ylim(top=250)"
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}