week12_nondiffnets/homework.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "IVAygknQX7jt"
   },
   "source": [
    "# REINFORCE in PyTorch\n",
    "\n",
    "При создании ноутбука использовались материалы курса [Practical RL](https://github.com/yandexdataschool/Practical_RL): [reference link](https://github.com/yandexdataschool/Practical_RL/blob/master/week06_policy_based/reinforce_pytorch.ipynb)\n",
    "\n",
    "В данном задании мы создадим модель для игры `CartPole-v0` с иcпользованием policy gradient (REINFORCE)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "axD5i_bFxRal"
   },
   "source": [
    "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/mryab/dl-hse-ami/blob/main/week12_nondiffnets/homework.ipynb)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "kNX2IP2UX7jt",
    "outputId": "81c4d3f3-a5c9-4347-8cfc-b3a69c2864b3"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Starting virtual X frame buffer: Xvfb.\n"
     ]
    }
   ],
   "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",
    "    !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": "STJi5A4zyPDL"
   },
   "outputs": [],
   "source": [
    "# install old version \n",
    "!pip install -q gym[classic_control]==0.15.3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "qIdPBHWWX7jt"
   },
   "outputs": [],
   "source": [
    "import gym\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "5eOABeuhX7jt"
   },
   "source": [
    "Потенциальная проблема: с некоторыми версиями `pyglet` следующая ячейка может упасть с ошибкой `NameError: name 'base' is not defined`. Ссылка на описание проблемы на GitHub находится [здесь](https://github.com/pyglet/pyglet/issues/134). Если Вы наблюдаете такую проблему, попробуйте перезапустить kernel."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "lJYyqd12X7ju"
   },
   "outputs": [],
   "source": [
    "env = gym.make(\"CartPole-v0\")\n",
    "\n",
    "# gym compatibility: unwrap TimeLimit\n",
    "if hasattr(env, '_max_episode_steps'):\n",
    "    env = env.env\n",
    "\n",
    "env.reset()\n",
    "n_actions = env.action_space.n\n",
    "state_dim = env.observation_space.shape\n",
    "\n",
    "plt.imshow(env.render(\"rgb_array\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "93ZrYvoZX7ju"
   },
   "source": [
    "# Создание сети для REINFORCE"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "ALC6Gt-MX7ju"
   },
   "source": [
    "Для алгоритма REINFORCE нам понадобится модель, которая предсказывает вероятность действий по состоянию среды.\n",
    "\n",
    "Для численной стабильности, пожалуйста, __не используйте softmax-слой в архитектуре модели__.\n",
    "Мы будем использовать softmax или log-softmax там, где это нужно.\n",
    "\n",
    "**Задание 1 (0.3 балла).** Реализуйте простую сеть-агента, принимающую на вход состояние среды и выдающую логиты для каждого действия."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "u72fskDGX7ju"
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "ap6y7QHWX7ju"
   },
   "outputs": [],
   "source": [
    "# Build a simple neural network that predicts policy logits. \n",
    "# Keep it simple: CartPole isn't worth deep architectures.\n",
    "model = nn.Sequential(\n",
    "  <YOUR CODE: define a neural network that predicts policy logits>\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "1mUBMieAX7ju"
   },
   "source": [
    "### Функция предсказания верояностей"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "x5gJjrc9X7ju"
   },
   "source": [
    "**Задание 2 (0.1 балла).**  Реализуйте функцию, выдающую вероятности действий при условии состояний.\n",
    "\n",
    "Заметка: выход этой функции не torch tensor, а numpy array. Поэтому нам не нужно считать в ней градиенты.\n",
    "<br>\n",
    "Используйте [no_grad](https://pytorch.org/docs/stable/autograd.html#torch.autograd.no_grad),\n",
    "чтобы отключить подсчет градиентов.\n",
    "<br>\n",
    "Также можно использовать `.detach()` (или `.data`), но эти два способа имеют следующие отличия:\n",
    "<br>\n",
    "\n",
    "*  При использовании `.detach()` граф вычислений строится, но затем отсоединяется от конкретного тензора, поэтому `.detach()` нужно использовать, если граф вычислений нужен нам для подсчета градиента через какой-то другой (не тот, на котором мы вызвали `.detach()`) тензор\n",
    "\n",
    "*   При использовании `no_grad()`, PyTorch не строит графа, поэтому в этой задаче данный способ предпочителен\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "e57rgiXvX7ju"
   },
   "outputs": [],
   "source": [
    "def predict_probs(states):\n",
    "    \"\"\" \n",
    "    Predict action probabilities given states.\n",
    "    :param states: numpy array of shape [batch, state_shape]\n",
    "    :returns: numpy array of shape [batch, n_actions]\n",
    "    \"\"\"\n",
    "    # convert states, compute logits, use softmax to get probability\n",
    "    <YOUR CODE>\n",
    "    return <YOUR CODE>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "mtUxAB9AX7ju"
   },
   "outputs": [],
   "source": [
    "test_states = np.array([env.reset() for _ in range(5)])\n",
    "test_probas = predict_probs(test_states)\n",
    "assert isinstance(test_probas, np.ndarray), \\\n",
    "    \"you must return np array and not %s\" % type(test_probas)\n",
    "assert tuple(test_probas.shape) == (test_states.shape[0], env.action_space.n), \\\n",
    "    \"wrong output shape: %s\" % np.shape(test_probas)\n",
    "assert np.allclose(np.sum(test_probas, axis=1), 1), \"probabilities do not sum to 1\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Eu6bRGizX7ju"
   },
   "source": [
    "### Сыграем в игру!\n",
    "\n",
    "Теперь мы можем использовать построенную модель, чтобы сыграть в игру!\n",
    "\n",
    "**Задание 3 (0.1 балла).** Заполните пропуски в коде одной сессии."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "Ji3FqW3WX7ju"
   },
   "outputs": [],
   "source": [
    "def generate_session(env, t_max=1000):\n",
    "    \"\"\" \n",
    "    Play a full session with REINFORCE agent.\n",
    "    Returns sequences of states, actions, and rewards.\n",
    "    \"\"\"\n",
    "    # arrays to record session\n",
    "    states, actions, rewards = [], [], []\n",
    "    s = env.reset()\n",
    "\n",
    "    for t in range(t_max):\n",
    "        # action probabilities array aka pi(a|s)\n",
    "        action_probs = predict_probs(np.array([s]))[0]\n",
    "\n",
    "        # Sample action with given probabilities.\n",
    "        a = <YOUR CODE>\n",
    "        new_s, r, done, info = env.step(a)\n",
    "\n",
    "        # record session history to train later\n",
    "        states.append(s)\n",
    "        actions.append(a)\n",
    "        rewards.append(r)\n",
    "\n",
    "        s = new_s\n",
    "        if done:\n",
    "            break\n",
    "\n",
    "    return states, actions, rewards"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "j68WQSOhX7ju"
   },
   "outputs": [],
   "source": [
    "# test it\n",
    "states, actions, rewards = generate_session(env)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "NyzxApMhX7ju"
   },
   "source": [
    "### Подсчет кумулятивной награды\n",
    "\n",
    "$$\n",
    "\\begin{align*}\n",
    "G_t &= r_t + \\gamma r_{t + 1} + \\gamma^2 r_{t + 2} + \\ldots \\\\\n",
    "&= \\sum_{i = t}^T \\gamma^{i - t} r_i \\\\\n",
    "&= r_t + \\gamma * G_{t + 1}\n",
    "\\end{align*}\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "t5R-ucSkBLXD"
   },
   "source": [
    "**Задание 4 (0.2 балла).** Реализуйте функцию, подсчитывающую кумулятивную награду с дисконтом $\\gamma$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "APWwH0yVX7ju"
   },
   "outputs": [],
   "source": [
    "def get_cumulative_rewards(rewards,  # rewards at each step\n",
    "                           gamma=0.99  # discount for reward\n",
    "                           ):\n",
    "    \"\"\"\n",
    "    Take a list of immediate rewards r(s,a) for the whole session \n",
    "    and compute cumulative returns (a.k.a. G(s,a) in Sutton '16).\n",
    "    \n",
    "    G_t = r_t + gamma*r_{t+1} + gamma^2*r_{t+2} + ...\n",
    "\n",
    "    A simple way to compute cumulative rewards is to iterate from the last\n",
    "    to the first timestep and compute G_t = r_t + gamma*G_{t+1} recurrently\n",
    "\n",
    "    You must return an array/list of cumulative rewards with as many elements as in the initial rewards.\n",
    "    \"\"\"\n",
    "    <YOUR CODE>\n",
    "    return <YOUR CODE: array of cumulative rewards>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "n-lLJb3UX7ju",
    "outputId": "3c90c87b-65cd-4226-945d-68b604dde9e5"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "looks good!\n"
     ]
    }
   ],
   "source": [
    "get_cumulative_rewards(rewards)\n",
    "assert len(get_cumulative_rewards(list(range(100)))) == 100\n",
    "assert np.allclose(\n",
    "    get_cumulative_rewards([0, 0, 1, 0, 0, 1, 0], gamma=0.9),\n",
    "    [1.40049, 1.5561, 1.729, 0.81, 0.9, 1.0, 0.0])\n",
    "assert np.allclose(\n",
    "    get_cumulative_rewards([0, 0, 1, -2, 3, -4, 0], gamma=0.5),\n",
    "    [0.0625, 0.125, 0.25, -1.5, 1.0, -4.0, 0.0])\n",
    "assert np.allclose(\n",
    "    get_cumulative_rewards([0, 0, 1, 2, 3, 4, 0], gamma=0),\n",
    "    [0, 0, 1, 2, 3, 4, 0])\n",
    "print(\"looks good!\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "q2uFzW6TX7ju"
   },
   "source": [
    "### Функция потерь и обучение\n",
    "\n",
    "Теперь нам нужно определить целевую функцию и обновление нашего policy gradient.\n",
    "\n",
    "Наша целевая функция задается вот так:\n",
    "\n",
    "$$ J \\approx  { 1 \\over N } \\sum_{s_i,a_i} G(s_i,a_i) $$\n",
    "\n",
    "REINFORCE дает нам способ подсчитать градиент математического ожидания награды по параметрам policy. Формула выглядит так:\n",
    "\n",
    "$$ \\nabla_\\theta \\hat J(\\theta) \\approx { 1 \\over N } \\sum_{s_i, a_i} \\nabla_\\theta \\log \\pi_\\theta (a_i \\mid s_i) \\cdot G_t(s_i, a_i) $$\n",
    "\n",
    "Мы можем использовать возможности PyTorch для автоматического дифференцирования, задав нашу целевую функцию таким образом:\n",
    "\n",
    "$$ \\hat J(\\theta) \\approx { 1 \\over N } \\sum_{s_i, a_i} \\log \\pi_\\theta (a_i \\mid s_i) \\cdot G_t(s_i, a_i) $$\n",
    "\n",
    "Когда мы подсчитаем градиент этой функции по отношению к весам $\\theta$, он станет в точности равен нашему policy gradient."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "a2DujhvbBaTi"
   },
   "source": [
    "**Задание 5 (0.3 балла).** Заполните пропуски в функции `train_on_session`, реализовав вычисление функции потерь и шаг градиентного спуска."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "aNFGvKp3X7ju"
   },
   "outputs": [],
   "source": [
    "def to_one_hot(y_tensor, ndims):\n",
    "    \"\"\" helper: take an integer vector and convert it to 1-hot matrix. \"\"\"\n",
    "    y_tensor = y_tensor.type(torch.LongTensor).view(-1, 1)\n",
    "    y_one_hot = torch.zeros(\n",
    "        y_tensor.size()[0], ndims).scatter_(1, y_tensor, 1)\n",
    "    return y_one_hot"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "ABaU0L5lX7jv"
   },
   "outputs": [],
   "source": [
    "# Your code: define optimizers\n",
    "optimizer = torch.optim.Adam(model.parameters(), 1e-3)\n",
    "\n",
    "\n",
    "def train_on_session(states, actions, rewards, gamma=0.99, entropy_coef=1e-2):\n",
    "    \"\"\"\n",
    "    Takes a sequence of states, actions and rewards produced by generate_session.\n",
    "    Updates agent's weights by following the policy gradient above.\n",
    "    Please use Adam optimizer with default parameters.\n",
    "    \"\"\"\n",
    "\n",
    "    # cast everything into torch tensors\n",
    "    states = torch.tensor(states, dtype=torch.float32)\n",
    "    actions = torch.tensor(actions, dtype=torch.int32)\n",
    "    cumulative_returns = np.array(get_cumulative_rewards(rewards, gamma))\n",
    "    cumulative_returns = torch.tensor(cumulative_returns, dtype=torch.float32)\n",
    "\n",
    "    # predict logits, probas and log-probas using an agent.\n",
    "    logits = model(states)\n",
    "    probs = nn.functional.softmax(logits, -1)\n",
    "    log_probs = nn.functional.log_softmax(logits, -1)\n",
    "\n",
    "    assert all(isinstance(v, torch.Tensor) for v in [logits, probs, log_probs]), \\\n",
    "        \"please use compute using torch tensors and don't use predict_probs function\"\n",
    "\n",
    "    # select log-probabilities for chosen actions, log pi(a_i|s_i)\n",
    "    log_probs_for_actions = torch.sum(\n",
    "        log_probs * to_one_hot(actions, env.action_space.n), dim=1)\n",
    "   \n",
    "    # Compute loss here. Don't forgen entropy regularization with `entropy_coef` \n",
    "    entropy = <YOUR CODE>\n",
    "    loss = <YOUR CODE>\n",
    "\n",
    "    # Gradient descent step\n",
    "    <YOUR CODE>\n",
    "\n",
    "    # technical: return session rewards to print them later\n",
    "    return np.sum(rewards)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Sc-w5BfSX7jv"
   },
   "source": [
    "### Обучение"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "Flcq3TCsX7jv"
   },
   "outputs": [],
   "source": [
    "for i in range(100):\n",
    "    rewards = [train_on_session(*generate_session(env)) for _ in range(100)]  # generate new sessions\n",
    "    \n",
    "    print(\"mean reward:%.3f\" % (np.mean(rewards)))\n",
    "    \n",
    "    if np.mean(rewards) > 500:\n",
    "        print(\"You Win!\")  # but you can train even further\n",
    "        break"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "6J-B06FfX7jv"
   },
   "source": [
    "### Результаты & видео"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "CyqIG1fPX7jv"
   },
   "outputs": [],
   "source": [
    "# Record sessions\n",
    "\n",
    "import gym.wrappers\n",
    "\n",
    "with gym.wrappers.Monitor(gym.make(\"CartPole-v0\"), directory=\"videos\", force=True) as env_monitor:\n",
    "    sessions = [generate_session(env_monitor) for _ in range(100)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "NUD7am78X7jv"
   },
   "outputs": [],
   "source": [
    "# Show video. This may not work in some setups. If it doesn't\n",
    "# work for you, you can download the videos and view them locally.\n",
    "# You may face 'No video with supported format and MIME type found' issue\n",
    "\n",
    "from pathlib import Path\n",
    "from IPython.display import HTML\n",
    "\n",
    "video_names = sorted([s for s in Path('videos').iterdir() if s.suffix == '.mp4'])\n",
    "\n",
    "HTML(\"\"\"\n",
    "<video width=\"640\" height=\"480\" controls>\n",
    "  <source src=\"{}\" type=\"video/mp4\">\n",
    "</video>\n",
    "\"\"\".format(video_names[-1]))  # You can also try other indices"
   ]
  }
 ],
 "metadata": {
  "colab": {
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.6"
  },
  "toc": {
   "base_numbering": 1,
   "nav_menu": {},
   "number_sections": true,
   "sideBar": true,
   "skip_h1_title": false,
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": false,
   "toc_position": {},
   "toc_section_display": true,
   "toc_window_display": false
  },
  "varInspector": {
   "cols": {
    "lenName": 16,
    "lenType": 16,
    "lenVar": 40
   },
   "kernels_config": {
    "python": {
     "delete_cmd_postfix": "",
     "delete_cmd_prefix": "del ",
     "library": "var_list.py",
     "varRefreshCmd": "print(var_dic_list())"
    },
    "r": {
     "delete_cmd_postfix": ") ",
     "delete_cmd_prefix": "rm(",
     "library": "var_list.r",
     "varRefreshCmd": "cat(var_dic_list()) "
    }
   },
   "types_to_exclude": [
    "module",
    "function",
    "builtin_function_or_method",
    "instance",
    "_Feature"
   ],
   "window_display": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}