Aula 4/Solução/Aula 4 - Parte Prática.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Aula 4 - Parte Prática - Actor-Critic (A2C)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Introdução\n",
    "\n",
    "Nesse quarto notebook vamos estudar um dos algoritmos mais implementados da família dos *Policy Gradients*. Ao implementar o A2C, você terá contato com importantes conceitos utilizados em *Deep RL*. Em particular, utilizaremos pela primeira vez no curso redes neurais com *features* compartilhadas e ambientes vetorizados (i.e., paralelizados) para coleta de dados mais eficiente.\n",
    "\n",
    "\n",
    "### Objetivos:\n",
    "\n",
    "- Familiarizar-se com os componentes *Actor* e *Critic*\n",
    "- Entender o papel da função Valor na estimativa truncada dos retornos\n",
    "- Ter um primeiro contato com truques de implementação tipicamente utilizados e RL\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Imports\n",
    "\n",
    "> **Atenção:** não se esqueça de executar todos os `imports` necessários antes prosseguir com o tutorial."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import logging\n",
    "from pprint import pprint\n",
    "\n",
    "import gym\n",
    "import numpy as np\n",
    "import tensorflow as tf\n",
    "import tensorflow_probability as tfp\n",
    "\n",
    "from utils.agent import RLAgent, RandomAgent\n",
    "from utils.memory import OnPolicyReplay\n",
    "from utils.networks import build_actor_critic_network\n",
    "import utils.runner\n",
    "from utils.viz import *\n",
    "\n",
    "\n",
    "tf.get_logger().setLevel(\"ERROR\")     # ignore TensorFlow warnings\n",
    "gym.logger.set_level(logging.ERROR)   # ignore OpenAI Gym warnings"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 0. LunarLander-v2\n",
    "\n",
    "Para o notebook de hoje utilizaremos um outro problema do Gym que é mais desafiador que o CartPole.\n",
    "\n",
    "\n",
    "> **Atenção**: para entender melhor a tarefa leia a documentação do LunarLander disponível em http://gym.openai.com/envs/LunarLander-v2/.\n",
    "\n",
    "Execute o código abaixo para visualizar alguns episódios do agente aleatório para ter uma melhor ideia da tarefa:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "env = gym.make(\"LunarLander-v2\")\n",
    "\n",
    "agent = RandomAgent(env.observation_space, env.action_space, None)\n",
    "\n",
    "utils.runner.evaluate(agent, env, n_episodes=10, render=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. Ambientes vetorizados no Gym\n",
    "\n",
    "Pela primeira vez no curso estaremos utilizando ambientes vetorizados, isto é, que emulam o comportamente de vários ambientes sendo executados em paralelo.\n",
    "\n",
    "Execute o código abaixo e tente entender como o ambiente retornado pelo `gym.vector.make` se comporta."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "env = gym.vector.make(\"LunarLander-v2\", num_envs=4, asynchronous=True)\n",
    "\n",
    "print(env.processes)\n",
    "\n",
    "print(env.observation_space, env.single_observation_space)\n",
    "print(env.action_space, env.single_action_space)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "n_steps_per_env = 3\n",
    "\n",
    "observations = env.reset()\n",
    "\n",
    "step = 0\n",
    "\n",
    "for _ in range(n_steps_per_env):\n",
    "    actions = env.action_space.sample()\n",
    "    observations, rewards, dones, _ = env.step(actions)\n",
    "    step += len(observations)\n",
    "\n",
    "    print(f\">> step = {step}\")\n",
    "    print(f\"observations =\\n{observations}\")\n",
    "    print(f\"actions = {actions}\")\n",
    "    print(f\"rewards = {rewards}\")\n",
    "    print(f\"dones = {dones}\")\n",
    "    print()\n",
    "    \n",
    "print(f\">> num_envs = {env.num_envs}, n_steps_per_env = {n_steps_per_env}, total timesteps = {step}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. Advantage Actor-Critic (A2C)\n",
    "\n",
    "<img src=\"img/a2c-algo.png\" alt=\"A2C Algorithm\" style=\"width: 700px;\"/>\n",
    "\n",
    "\n",
    "> **Observação**: para uma introdução mais intuitiva do A2C recomendamos o blog post https://sudonull.com/post/32170-Intuitive-RL-Reinforcement-Learning-Introduction-to-Advantage-Actor-Critic-A2C. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2.1 Compartilhando informação entre *Actor* e *Critic*: *2-head model*\n",
    "\n",
    "Na arquitetura A2C é comum implementar o Actor e o Critic em um mesmo modelo que compartilha parâmetros. Execute o código abaixo e tente interpretar a figura."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "num_envs = 8\n",
    "env = gym.vector.make(\"LunarLander-v2\", num_envs=num_envs, asynchronous=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "config = {\n",
    "    \"actor_critic_net\": {\n",
    "        \"hidden_layers\": [64, 64],\n",
    "        \"activation\": \"tanh\"\n",
    "    },\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "actor_critic_net = build_actor_critic_network(env.single_observation_space, env.single_action_space, config[\"actor_critic_net\"])\n",
    "tf.keras.utils.plot_model(actor_critic_net, show_shapes=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2.2 Otimizando a política e a função Valor via *joint loss*\n",
    "\n",
    "$$\n",
    "[\\theta, \\phi] \\leftarrow [\\theta, \\phi] + \\alpha \\nabla_{\\theta, \\phi}(L_{actor}(\\theta) + L_{critic}(\\phi))\n",
    "$$\n",
    "\n",
    "onde $L_{actor}(\\theta)$ e $L_{critic}(\\phi)$ correspondem respectivamente ao *policy loss* e *mean squared error*:\n",
    "$$\n",
    "\\begin{align*}\n",
    "L_{actor}(\\theta) &= - \\frac{1}{K} \\sum_{t=1}^K \\log \\pi_{\\theta}(\\mathbf{a}_t|\\mathbf{s}_t) \\hat{A}_t^{(n)} \\\\\n",
    "L_{critic}(\\phi) &= \\frac{1}{K} \\sum_{t=1}^K  (V_{\\phi}(\\mathbf{s}_t) - \\hat{R}_t)^2\n",
    "\\end{align*}\n",
    "$$\n",
    "\n",
    "> **<font color=\"red\">IMPORTANTE</color>**: Note que para o problema de regressão que precisamos resolver para aprender o *critic*, o retorno descontado $\\hat{R}_t$ funciona como o *target* e, portanto, deve ser considerado uma \"constante\" para o TensorFlow. No exercício abaixo você precisará se lembrar disso!\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2.3  Bônus de exploração: entropia da distribuição de ações\n",
    "\n",
    "Além do *joint loss* adicionaremos também um outro termo à função objetivo a fim de incentivar o agente a continuar explorando novas ações.\n",
    "\n",
    "$$\n",
    "H(\\pi_{\\theta}) = \\mathbb{E}_{\\mathbf{s}, \\mathbf{a} \\sim \\sim \\pi_{\\theta}} \\left[ - \\log \\pi_{\\theta}(\\mathbf{a}|\\mathbf{s}) \\right]\n",
    "$$\n",
    "\n",
    "Tente entender no código abaixo onde e como esse bônus da entropia é implementado."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. Detalhes (importantes) de implementação\n",
    "\n",
    "### 3.1 Truncando a norma dos gradientes\n",
    "\n",
    "$$\n",
    "\\tilde\\nabla L(\\theta) = \n",
    "\\begin{cases}\n",
    "    \\nabla L(\\theta)                                     & \\text{ se } \\| \\nabla L(\\theta) \\| < clip \\\\\n",
    "    \\frac{\\nabla L(\\theta)}{\\| \\nabla L(\\theta) \\|} clip & \\text{ c.c. }\n",
    "\\end{cases}\n",
    "$$\n",
    "\n",
    "No código da classe abaixo, implementamos um truque conhecido em Deep Learning como *gradient clipping*. Ele consiste em preprocessar os gradientes de cada peso da rede truncando a sua magnitude caso ela passe de um certo limiar, ilustrado acima com o gradiente de uma função-objetivo $L(\\theta)$ genérica. Esse truque é comum pois gradientes podem tomar magnitudes muito grandes quando o ambiente é ruidoso ou a rede muito grande. Isso gera instabilidades no treinamento e pode até levar pesos a valores inválidos como *Not a Number* (`Nan`). Implementamos esse preprocessamento logo antes de chamar `apply_gradients` no método `learn()` do agente.\n",
    "\n",
    "### 3.2 *Learning rate scheduler*\n",
    "\n",
    "$$\n",
    "\\theta \\gets \\theta - \\alpha_t \\nabla L(\\theta)\n",
    "$$\n",
    "\n",
    "O segundo truque que implementamos abaixo é o uso de um agendador para a taxa de aprendizado (`learning rate`) $\\alpha_t$. A equação acima ilustra a regra genérica de atualização dos parâmetros $\\theta$ para alguma função objetivo $L(\\theta)$. Por mais que otimizadores como Adam e RMSprop já façam uma escolha bem motivada de $\\alpha_t$ de acordo com o valor base $\\alpha$ passado ao construtor, foi observado empiricamente que diminuir $\\alpha$ ao longo do treinamento ajuda a estabilizar o treinamento. Implementamos isso com o uso do `tf.keras.optimizers.schedules.PolynomialDecay` no lugar de um valor fixo para o `learning_rate` do RMSprop. Configuramos o *scheduler* de forma que $\\alpha_t$ decresca linearmente até o final do aprendizado.\n",
    "\n",
    "Intuitivamente, imagine que a política fique mais sensível a variações nos parâmetros $\\theta$ com o passar do tempo. Para observar esse fenômeno, experimente passar um valor fixo para o `learning_rate` abaixo. É comum observar que, após um período inicial de melhora nas curvas de retorno acumulado, o desempenho cai (muitas vezes abruptamente) perto do final do treinamento.\n",
    "\n",
    "### 3.3 Ponderando os componentes do *joint loss*\n",
    "\n",
    "$$\n",
    "L(\\theta, \\phi) = L_\\text{actor} (\\theta) + \\beta L_\\text{critic}(\\phi) - \\epsilon H(\\phi)\n",
    "$$\n",
    "\n",
    "Outro aspecto crucial de implementação é a ponderação adequada da contribuição das *losses* de cada componente para o *joint loss*. De fato, é raro que as funções objetivo de *actor*, *critic* e entropia estejam na mesma escala: muitas vezes $L_\\text{actor}(\\theta)$ assume valores pequenos enquanto $L_\\text{critic}(\\phi)$, valores grandes, principalmente no início do treinamento quando a função-valor está mal-ajustada.\n",
    "\n",
    "É fundamental então experimentar com coeficientes diferentes para os objetivos do *critic* e de entropia, implementados abaixo, respectivamente, com os campos `vf_loss_coeff` e `entropy_coeff` do `config`. Os valores padrão já foram ajustados, mas experimente mudá-los para atingir um desempenho melhor na tarefa."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "\n",
    "**<font color=\"red\">EXERCÍCIO-PROGRAMA:</font>**\n",
    "\n",
    "Nesse exercício você deverá implementar o *value function loss* no método `_joint_loss_fn` da classe abaixo.\n",
    "\n",
    "\n",
    "> **Nota 1**: consulte a documentação do MSE em https://www.tensorflow.org/api_docs/python/tf/keras/losses/MeanSquaredError. Será útil!\n",
    "\n",
    "> **Nota 2**: você precisará utilizar a função `tf.stop_gradient` no cálculo do *target* na loss do Value Function. Isso permite que o TensorFlow não tente diferenciar o *target*. Veja nota <font color=\"red\">IMPORTANTE</color> na Seção 2.2.\n",
    "\n",
    "> **Nota 3**: se você tiver algum erro de \"tipos\", por exemplo, esperava-se `float32` em algum ponto do código mas o array do *NumPy* foi calculado como `float64` você poderá fazer *casting* manual de tipos usando o método de `astype(\"f\")`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class A2C(RLAgent):\n",
    "    \n",
    "    def __init__(self, obs_space, action_space, config):\n",
    "        super().__init__(obs_space, action_space, config)\n",
    "        \n",
    "        self.memory = OnPolicyReplay()\n",
    "        self.actor_critic = build_actor_critic_network(obs_space, action_space, config[\"actor_critic_net\"])\n",
    "\n",
    "        self.optimizer = tf.keras.optimizers.RMSprop(\n",
    "            learning_rate=tf.keras.optimizers.schedules.PolynomialDecay(\n",
    "                0.00083,\n",
    "                decay_steps=config[\"total_timesteps\"] / config[\"train_batch_size\"],\n",
    "                end_learning_rate=1e-4,\n",
    "                power=1.0\n",
    "            ))\n",
    "\n",
    "    def act(self, obs):\n",
    "        \"\"\"\n",
    "        Escolhe uma ação para ser tomada dada uma observação do ambiente.\n",
    "        \n",
    "        Args: \n",
    "            obs: observação do ambiente.\n",
    "        \n",
    "        Return:\n",
    "            action: ação válida dentro do espaço de ações.\n",
    "        \"\"\"\n",
    "        return self._act(obs).numpy()\n",
    "        \n",
    "    @tf.function\n",
    "    def _act(self, obs):\n",
    "        action_dist, _ = self.actor_critic(obs)\n",
    "        return action_dist.sample()\n",
    "    \n",
    "    def observe(self, obs, action, reward, next_obs, done):\n",
    "        \"\"\"\n",
    "        Registra na memória do agente uma transição do ambiente.\n",
    "\n",
    "        Args:\n",
    "            obs:            observação do ambiente antes da execução da ação.\n",
    "            action:         ação escolhida pelo agente.\n",
    "            reward (float): escalar indicando a recompensa obtida após a execução da ação.\n",
    "            next_obs:       nova observação recebida do ambiente após a execução da ação.\n",
    "            done (bool):    True se a nova observação corresponde a um estado terminal, False caso contrário.\n",
    "\n",
    "        Return:\n",
    "            None\n",
    "        \"\"\"\n",
    "        self.memory.update(obs, action, reward, next_obs, done)\n",
    "\n",
    "    def learn(self):\n",
    "        \"\"\"\n",
    "        Método de treinamento do agente. A partir das experiências de sua memória,\n",
    "        o agente aprende um novo comportamento.\n",
    "\n",
    "        Args: \n",
    "            None\n",
    "\n",
    "        Return:\n",
    "            None\n",
    "        \"\"\"\n",
    "        if self.memory.batch_size < self.config[\"train_batch_size\"]:\n",
    "            return\n",
    "        \n",
    "        batch = self.memory.sample()\n",
    "        weights = self.actor_critic.trainable_weights\n",
    "\n",
    "        with tf.GradientTape() as tape:\n",
    "            policy_loss, vf_loss, entropy_loss = self._joint_loss_fn(batch)\n",
    "            loss = policy_loss + self.config[\"vf_loss_coeff\"] * tf.cast(vf_loss, tf.float32) - self.config[\"entropy_coeff\"] * entropy_loss\n",
    "            gradients = tape.gradient(loss, weights)\n",
    "    \n",
    "        gradients = tuple(tf.clip_by_norm(grad, clip_norm=0.5) for grad in gradients)\n",
    "        self.optimizer.apply_gradients(zip(gradients, weights))\n",
    "      \n",
    "        return {\n",
    "            \"policy_loss\": policy_loss.numpy(),\n",
    "            \"vf_loss\": vf_loss.numpy(),\n",
    "            \"entropy_loss\": entropy_loss.numpy()\n",
    "        }\n",
    "\n",
    "    def _joint_loss_fn(self, batch):\n",
    "        \"\"\"\n",
    "        Calcula a função loss do policy gradients para um `batch` de transições.\n",
    "        \n",
    "        Um `batch` agrega arrays n-dimensionais. Cada array (e.g., batch[\"states\"],\n",
    "        batch[\"actions\"], batch[\"rewards\"]) tem como primeiras duas dimensões o número\n",
    "        de passos dados no ambiente vetorizado e o número de ambientes em paralelo. \n",
    "        Por exemplo, batch[\"states\"][t][k] devolve um array correspondendo ao estado \n",
    "        no passo t devolvido pelo k-ésimo ambiente.\n",
    "\n",
    "        Args:\n",
    "            batch (Dict[str, np.ndarray]): dicionário para acesso às matrizes de \n",
    "                estados, ações, recompensas, próximos estados e flags de terminação. \n",
    "        \n",
    "        Return:\n",
    "            loss (tf.Tensor): surrogate loss conjunta da política, função valor e\n",
    "                bônus de entropia.\n",
    "        \"\"\"\n",
    "        states = batch[\"states\"]\n",
    "        actions = batch[\"actions\"]\n",
    "        rewards = batch[\"rewards\"]\n",
    "        next_states = batch[\"next_states\"]\n",
    "        dones = batch[\"dones\"]\n",
    "\n",
    "        n_steps = len(states)\n",
    "        gamma = self.config[\"gamma\"]\n",
    "        lambda_ = self.config[\"lambda\"]\n",
    "\n",
    "        action_dists, values = self.actor_critic(states)\n",
    "        _, last_value = self.actor_critic(next_states[-1:])\n",
    "\n",
    "        values = tf.squeeze(tf.concat([values, last_value], axis=0))\n",
    "        values, next_values = values[:-1], values[1:]\n",
    "\n",
    "        deltas = rewards + gamma * (1 - dones) * next_values - values\n",
    "\n",
    "        returns = np.empty_like(rewards)\n",
    "        advantages = np.empty_like(rewards)\n",
    "\n",
    "        returns[-1] = rewards[-1] + gamma * (1 - dones[-1]) * next_values[-1]\n",
    "        advantages[-1] = deltas[-1]\n",
    "\n",
    "        for t in reversed(range(n_steps - 1)):\n",
    "            returns[t] = rewards[t] + gamma * (1 - dones[t]) * returns[t+1]\n",
    "            advantages[t] = deltas[t] + (gamma * lambda_) * (1 - dones[t]) * advantages[t+1]\n",
    "\n",
    "        log_probs = action_dists.log_prob(actions)\n",
    "\n",
    "        policy_loss = - tf.reduce_sum(log_probs * tf.stop_gradient(advantages.astype(\"f\")))\n",
    "        \n",
    "        # SEU CÓDIGO AQUI =====================================\n",
    "        vf_loss = None\n",
    "        \n",
    "    \n",
    "        # =====================================================\n",
    "        entropy_loss = tf.reduce_mean(action_dists.entropy())\n",
    "\n",
    "        return policy_loss, vf_loss, entropy_loss\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Para treinar o seu agente A2C execute o código abaixo:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "total_timesteps = 500_000\n",
    "\n",
    "num_envs = 8\n",
    "train_env = gym.vector.make(\"LunarLander-v2\", num_envs=num_envs, asynchronous=True)\n",
    "\n",
    "config = {\n",
    "    \"actor_critic_net\": {\n",
    "        \"hidden_layers\": [64, 64],\n",
    "        \"activation\": \"tanh\"\n",
    "    },\n",
    "    \"optimizer\": {\n",
    "        \"class_name\": \"RMSprop\",\n",
    "        \"config\": {\n",
    "            \"learning_rate\": 8e-4,\n",
    "            \"rho\": 0.99\n",
    "        }\n",
    "    },\n",
    "    \"total_timesteps\": total_timesteps,\n",
    "    \"train_batch_size\": 40,\n",
    "    \"gamma\": 0.995,\n",
    "    \"lambda\": 1.0,\n",
    "    \"vf_loss_coeff\": 0.25,\n",
    "    \"entropy_coeff\": 1e-5\n",
    "}\n",
    "\n",
    "agent = A2C(train_env.single_observation_space, train_env.single_action_space, config)\n",
    "\n",
    "timesteps, avg_total_rewards, losses = utils.runner.train(agent, train_env, total_timesteps)\n",
    "\n",
    "plot_metrics(avg_total_rewards, losses)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Para visualizar o agente treinado execute o código abaixo:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "eval_env = gym.make(\"LunarLander-v2\")\n",
    "utils.runner.evaluate(agent, eval_env, n_episodes=10, render=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "\n",
    "**<font color=\"red\">PARABÉNS!</font>**\n",
    "\n",
    "\n",
    "Se você conseguiu chegar até aqui, parabéns! Você conseguiu fazer funcionar uma primeira versão do A2C. Sabemos que isso não é uma tarefa fácil... :)\n",
    "\n",
    "Esperamos que você tenha se familiarizado com as principais ideias por trás de um algoritmo de Deep RL e tenha sentido um pouco na pele a dificuladade de treinar agentes de RL, mesmo para problemas aparentemente simples."
   ]
  }
 ],
 "metadata": {
  "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"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}