04_curiosity.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "flatland-rl-workshop-2021-wwwconf-curiosity.ipynb",
      "provenance": [],
      "collapsed_sections": [],
      "toc_visible": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "oj7aWjVf4hXJ"
      },
      "source": [
        "Google Colab Setup\n",
        "---\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "PGLxbC7J4ver"
      },
      "source": [
        "#@title << Setup Google Colab by running this cell {display-mode: \"form\"}\n",
        "import sys\n",
        "if 'google.colab' in sys.modules:\n",
        "    # Clone GitHub repository\n",
        "    !git clone --single-branch --branch flatland-model-evaluator https://github.com/YanickSchraner/rl-on-trains-workshop.git\n",
        "        \n",
        "    # Copy files required to run the code\n",
        "    !cp -r \"rl-on-trains-workshop/utils\" \"rl-on-trains-workshop/train.py\" \"rl-on-trains-workshop/tests.yaml\" .\n",
        "    \n",
        "    # Install packages via pip\n",
        "    !pip install flatland-rl==2.2.2  aicrowd-cli\n",
        "    \n",
        "    # Restart Runtime\n",
        "    import os\n",
        "    os.kill(os.getpid(), 9)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "obSV5Qsx51_c"
      },
      "source": [
        "Tensorboard\n",
        "---\n",
        "\n",
        "Setup tensorboard before starting the training."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Q2iT6QbAFMS5"
      },
      "source": [
        "\"\"\" load tensorboard extension\"\"\"\n",
        "%load_ext tensorboard"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "tK_-M24-FOmL"
      },
      "source": [
        "\"\"\" startup tensorboad \"\"\"\n",
        "%tensorboard --logdir runs"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "L9k2qzoP9iIV"
      },
      "source": [
        "Random Network Distillation\n",
        "---\n",
        "Train a Flatland-RL policy using DQN and Random Network Distillation"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "o_q-yDOmm6li"
      },
      "source": [
        "import os\n",
        "from argparse import Namespace\n",
        "from pprint import pprint\n",
        "\n",
        "from train import train_agent\n",
        "from utils.dddqn_rnd_policy import DDDQNRNDPolicy"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "PD49_sSw-ZY-"
      },
      "source": [
        "Set policy class to the DDQN RND policy that combines DQN and Random Network Distillation"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "bqxpCt1B-W9D"
      },
      "source": [
        "policy_class = DDDQNRNDPolicy"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "WZtC_oUZ9qQ1"
      },
      "source": [
        "Training Parameters"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "tnEj2PoC9c-w"
      },
      "source": [
        "training_params = {\n",
        "    'n_episodes': 500,\n",
        "    'n_evaluation_episodes': 50,\n",
        "    'checkpoint_interval': 100,\n",
        "    'eps_start': 1.0,\n",
        "    'eps_end': 0.01,\n",
        "    'eps_decay': 0.9,\n",
        "    'buffer_size': int(1e5),\n",
        "    'buffer_min_size': 0,\n",
        "    'batch_size': 128,\n",
        "    'gamma': 0.95,\n",
        "    'tau': 1e-3,\n",
        "    'learning_rate': 1e-4,\n",
        "    'hidden_layers': [32, 32],\n",
        "    'update_every': 8,\n",
        "    'use_gpu': True,\n",
        "    'num_threads': 1,\n",
        "    'render': False, # won't work in notebook\n",
        "    'load_policy': '',\n",
        "    \n",
        "    # RND parameters:\n",
        "    'rnd_hidden_layers': [32, 32],\n",
        "    'rnd_intrinsic_reward_weight': 1.0,\n",
        "}"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gz-PSuoL9-is"
      },
      "source": [
        "Environment Parameters"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "wp82KbFn99zp"
      },
      "source": [
        "training_env_params = {\n",
        "    \"n_agents\": 5,\n",
        "    \"x_dim\": 25,\n",
        "    \"y_dim\": 25,\n",
        "    \"n_cities\": 2,\n",
        "    \"max_rails_between_cities\": 2,\n",
        "    \"max_rails_in_city\": 3,\n",
        "    \"malfunction_rate\": 1 / 50,\n",
        "    \"seed\": 0\n",
        "}\n",
        "evaluation_env_params = training_env_params"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "K1N5cYeT-GwA"
      },
      "source": [
        "Observation Parameters"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "dxcigYk_-I-s"
      },
      "source": [
        "obs_params = {\n",
        "    \"use_fast_tree_observation\": True,\n",
        "    \"observation_tree_depth\": 2,\n",
        "    \"observation_radius\": 10,\n",
        "    \"observation_max_path_depth\": 30\n",
        "}"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ploxLalg-RCx"
      },
      "source": [
        "Run Training"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "pAHfIF9D-QBL"
      },
      "source": [
        "print(\"\\nTraining parameters:\")\n",
        "pprint(training_params)\n",
        "\n",
        "print(\"\\nTraining environment parameters:\")\n",
        "pprint(training_env_params)\n",
        "\n",
        "print(\"\\nEvaluation environment parameters:\")\n",
        "pprint(evaluation_env_params)\n",
        "\n",
        "print(\"\\nObservation parameters:\")\n",
        "pprint(obs_params)\n",
        "\n",
        "os.environ[\"OMP_NUM_THREADS\"] = str(training_params['num_threads'])\n",
        "\n",
        "train_agent(policy_class,\n",
        "            Namespace(**training_params),\n",
        "            Namespace(**training_env_params),\n",
        "            Namespace(**evaluation_env_params),\n",
        "            Namespace(**obs_params))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "x__NbTEedNgz"
      },
      "source": [
        "Submit Agent\n",
        "---"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "iYH-ZA3Fdal-"
      },
      "source": [
        "!aicrowd login --api-key $API_KEY\n",
        "!aicrowd submission create -c flatland-workshop -f \"runs/2021-04-12_18-46-22/ckpts/ckpt-500.pth\""
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}