Aula 3/Aula-3/Aula 3 - Parte Prática - Redução de Variância e Função Valor.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Aula 3 - Parte Prática - Redução de Variância e Função Valor"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Introdução\n",
    "\n",
    "Nesse terceiro notebook vamos realizar experimentos com algumas técnicas estatísticas para redução de variância do estimador do *policy gradients* do algoritmo REINFORCE visto na última aula.\n",
    "\n",
    "\n",
    "$$\n",
    "\\nabla_\\theta J(\\theta) = \\mathbb{E}_{\\tau \\sim \\pi_\\theta} \\left[ \\sum_{t=0}^{T-1} \\nabla_\\theta \\log \\pi_\\theta(\\mathbf{a}_t|\\mathbf{s}_t) \\left ( \\left( \\sum_{k=t}^{T-1} r_{k} \\right)  - b(\\mathbf{s}_{t}) \\right ) \\right]\n",
    "$$\n",
    "\n",
    "\n",
    "### Objetivos:\n",
    "\n",
    "- Relacionar as propriedades do estimador REINFORCE com a performance do agente\n",
    "- Verificar experimentalmente o efeito de redução de variância do estimador de Policy Gradient calculado com *reward-to-go*\n",
    "- Incorporar a função Valor como *baseline* para os retornos das trajetórias no REINFORCE\n",
    "- Familiarizar-se com o aprendizado de função Valor via regressão sobre os retornos das trajetórias "
   ]
  },
  {
   "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",
    "\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\n",
    "from utils.memory import OnPolicyReplay\n",
    "from utils.networks import build_discrete_policy, build_value_network, get_optimizer\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. Configurações "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> **Atenção**: a fim de tornar o problema do `CartPole` um pouco mais desafiador, vamos utilizar nesse notebook a versão `v1` que aumenta o tamanho máximo de uma trajetória (i.e., `env.spec.max_episode_steps`) e também o retorno mínimo para resolver a tarefa (e.g, `env.spec.reward_threshold`)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "env = gym.make(\"CartPole-v1\")\n",
    "print(env.spec.max_episode_steps, env.spec.reward_threshold)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "config = {\n",
    "    # policy network\n",
    "    \"hidden_layers\": [64, 64],\n",
    "    \"activation\": \"relu\",\n",
    "\n",
    "    # optimization\n",
    "    \"optimizer\": \"adam\",\n",
    "    \"learning_rate\": 1e-3,\n",
    "\n",
    "    # training\n",
    "    \"train_batch_size\": 3000,\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. REINFORCE\n",
    "\n",
    "Na classe REINFORCE abaixo re-implementamos o policy gradient que estudamos na última aula e adicionamos um argumento adicional ao construtor `__init__` da classe correspondente a função `postprocessing` que processa um `batch` de recompensas. Utilizaremos esse argumento na próxima seção para passar uma função que computará o *reward-to-go* dos passos de uma trajetória."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class REINFORCE(RLAgent):\n",
    "    \"\"\"\n",
    "    Classe que implementa os componentes de um agente que aprende pelo REINFORCE.\n",
    "\n",
    "    Args:\n",
    "        obs_space:       especificação do espaço de observações do ambiente.\n",
    "        action_space:    especificação do espaço de ações do ambiente.\n",
    "        config (dict):   configurações de hiper-parâmetros.\n",
    "        postprocessing:  função que processa um batch de recompensas.\n",
    "    \"\"\"\n",
    "    \n",
    "    def __init__(self, obs_space, action_space, config, postprocessing=None):\n",
    "        super(REINFORCE, self).__init__(obs_space, action_space, config)\n",
    "        \n",
    "        self.memory = OnPolicyReplay()\n",
    "        self.policy = build_discrete_policy(self.obs_space, self.action_space, config[\"hidden_layers\"], config[\"activation\"])\n",
    "        self.optimizer = get_optimizer(config[\"optimizer\"], config[\"learning_rate\"])\n",
    "\n",
    "        self.postprocessing = postprocessing or np.sum\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.policy(obs[None,:])\n",
    "        return action_dist.sample()[0]\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",
    "\n",
    "        with tf.GradientTape() as tape:\n",
    "            loss = self._loss_pg_fn(batch)\n",
    "            gradients = tape.gradient(loss, self.policy.trainable_weights)\n",
    "\n",
    "        self.optimizer.apply_gradients(zip(gradients, self.policy.trainable_weights))\n",
    "      \n",
    "        return loss\n",
    "\n",
    "    def _loss_pg_fn(self, batch):\n",
    "        \"\"\"\n",
    "        Calcula a função loss do policy gradients para um `batch` de trajetórias/episódios.\n",
    "        \n",
    "        Um `batch` agrega listas de arrays n-dimensionais. Cada lista (e.g., batch[\"states\"],\n",
    "        batch[\"actions\"], batch[\"rewards\"]) tem o tamanho do número de episódios. Por exemplo,\n",
    "        batch[\"states\"][k] devolve um array n-dimensional para o k-ésimo episódio. Este array\n",
    "        tem como primeira dimensão o número de timesteps do k-ésimo episódio.\n",
    "\n",
    "        Args:\n",
    "            batch (Dict[str, List[np.ndarray]]): dicionário para acesso às listas de estados, ações e recompensas. \n",
    "        \n",
    "        Return:\n",
    "            loss (tf.Tensor): média sobre os episódios do surrogate loss function.\n",
    "            \n",
    "        \"\"\"\n",
    "        states, actions, rewards = batch[\"states\"], batch[\"actions\"], batch[\"rewards\"]\n",
    "        \n",
    "        n_episodes = len(rewards)\n",
    "\n",
    "        loss = 0.0\n",
    "        for episode in range(n_episodes):\n",
    "            action_dist = self.policy(states[episode])\n",
    "            log_prob = action_dist.log_prob(actions[episode])\n",
    "            R_t = self.postprocessing(rewards[episode])\n",
    "            loss += - tf.reduce_sum(log_prob * R_t)\n",
    "\n",
    "        loss /= n_episodes\n",
    "            \n",
    "        return loss"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Execute o código abaixo rodar o mesmo experimento do **REINFORCE** um número dado de vezes (i.e., `n_trials`):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "n_trials = 5\n",
    "total_timesteps = 1_000_000\n",
    "agent_cls = REINFORCE\n",
    "postprocessing = None\n",
    "\n",
    "timesteps, total_rewards, avg_total_rewards = utils.runner.run_experiments(\n",
    "    n_trials, env, agent_cls, config, postprocessing, total_timesteps)\n",
    "\n",
    "plot_experiments(env, timesteps, total_rewards, avg_total_rewards)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Execute o código abaixo para visualizar a política aprendida pelo algoritmo REINFORCE: "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "n_episodes = 10\n",
    "_ = evaluate(agent, env, n_episodes, render=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. Policy Gradients: ignorando recompensas passadas"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Inicialmente vamos substituir o retorno (i.e., recompensa total acumulada) de um episódio pelo *reward-to-go*:\n",
    "$$\n",
    "R_t = \\sum_{k=t}^{T-1} r_k\n",
    "$$\n",
    "\n",
    "Lembre-se que na fórmula original do *policy gradient* o retorno da trajetória inteira é utilizado como sinal de reforço para todos os passos $t$ da mesma forma. Em outras palavras, o *score function* $\\nabla_\\theta \\log \\pi_\\theta(\\mathbf{a}_t|\\mathbf{s}_t)$ do passo $t$ é ponderado com recompensas tanto do passado ($t' < t$) quanto do futuro ($t' \\geq t$)! Isso apenas adiciona ruído no estimador; impactando de forma negativa a sua variância. É importante notar que a escolha de uma ação só tem influência no retorno futuro, isto é, a partir do momento da tomada da ação.\n",
    "\n",
    "Ignorando recompensas passadas no REINFORCE podemos \"filtrar\" esse ruído de maneira que o estimador do policy gradient derivado anteriormente terá sua variância reduzida (embora seu valor esperado não se altere):\n",
    "\n",
    "$$\n",
    "\\begin{align*}\n",
    "\\nabla_\\theta J(\\theta) &= \\mathbb{E}_{\\tau \\sim \\pi_\\theta} \\left [ \\sum_{t=0}^T \\nabla_\\theta \\log \\pi_\\theta(\\mathbf{a}_t|\\mathbf{s}_t) \\left( \\sum_{t=1}^{T-1} r_t \\right) \\right ] \\\\\n",
    "&= \\mathbb{E}_{\\tau \\sim \\pi_\\theta} \\left [ \\sum_{t=0}^T \\nabla_\\theta \\log \\pi_\\theta(\\mathbf{a}_t|\\mathbf{s}_t) \\left( \\sum_{k=t}^{T-1} r_k \\right) \\right ]\n",
    "\\end{align*}\n",
    "$$\n",
    "\n",
    "> **Atenção**: note como a mudança do subescrito $k$ faz diferença na última igualdade! Na primeira igualdade (i.e.,"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2.1 Reward-to-Go \n",
    "\n",
    "No código abaixo implementamos o *reward-to-go* para uma trajetória. Note que a função do `np.cumsum` do NumPy calcula a soma acumulada a cada passo do array n-dimensional de `rewards`. Para maiores detalhes consulte a documentação de [numpy.cumsum](https://docs.scipy.org/doc/numpy/reference/generated/numpy.cumsum.html)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def compute_reward_to_go(rewards):\n",
    "#     returns = []\n",
    "#     total_reward = 0.0\n",
    "\n",
    "#     for i in range(len(rewards) - 1, -1, -1):\n",
    "#         reward = rewards[i]\n",
    "#         total_reward += reward\n",
    "#         returns.append(total_reward)\n",
    "\n",
    "#     returns = returns[::-1]\n",
    "\n",
    "#     return np.array(returns)\n",
    "    return np.cumsum(rewards[::-1])[::-1] # essa linha é equivalente ao for-loop comentado acima, porém mais eficiente!\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Execute o código abaixo rodar o mesmo experimento do **REINFORCE+reward-to-go** um número dado de vezes (i.e., `n_trials`):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "n_trials = 5\n",
    "total_timesteps = 1_000_000\n",
    "agent_cls = REINFORCE\n",
    "postprocessing = compute_reward_to_go\n",
    "\n",
    "timesteps, total_rewards, avg_total_rewards = utils.runner.run_experiments(\n",
    "    n_trials, env, agent_cls, config, postprocessing, total_timesteps)\n",
    "\n",
    "plot_experiments(env, timesteps, total_rewards, avg_total_rewards)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Execute o código abaixo para visualizar a política aprendida pelo algoritmo **REINFORCE + reward-to-go**: "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "n_episodes = 10\n",
    "_ = evaluate(agent, env, n_episodes, render=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. Policy Gradients: adicionando referências para os retornos (baseline)\n",
    "\n",
    "Na classe `VPG` (i.e., *Vanilla Policy Gradients*) especializamos a classe do algoritmo REINFORCE para incorporar uma referência para os retornos dos episódios (i.e., *baseline*). Neste notebook utilizaremos a função Valor como função *baseline*.\n",
    "\n",
    "\n",
    "> **Observação**: embora não haja consenso na literatura, note que estamos usando o nome de *Vanilla Policy Gradients* para denotar o algoritmo REINFORCE com *baseline* calculado pela função Valor.\n",
    "\n",
    "\n",
    "Lembre-se que a a **função Valor** tenta estimar o retorno médio a partir de um dado estado $\\mathbf{s}$:\n",
    "$$\n",
    "V_\\phi(\\mathbf{s}) \\approx V^{\\pi_\\theta}(\\mathbf{s}) = \\mathbb{E}_{\\tau \\sim \\pi_\\theta} \\left [ \\sum_{t=1}^{T-1} r_t~\\middle |~\\mathbf{s}_0 = \\mathbf{s} \\right ]\n",
    "$$\n",
    "\n",
    "Para aprendermos os parâmetros $\\phi$ do aproximador da função Valor tentaremos resolver um problema de [regressão](https://en.wikipedia.org/wiki/Nonlinear_regression) usando MSE (i.e., *[Mean Squared Error](https://en.wikipedia.org/wiki/Mean_squared_error)*) em cada época de treinamento (e.g., cada vez que o método `learn` é chamado dentro do ciclo de interação agente-ambiente):\n",
    "\n",
    "$$\n",
    "\\phi_k = \\arg\\min_{\\phi} \\mathbb{E}_{\\mathbf{s} \\sim d^{\\pi_{\\theta}}} \\left [ (V_{\\phi}(\\mathbf{s_t}) - R_t)^2 \\right ]\n",
    "$$\n",
    "onde $R_t = \\sum_{k=t}^{T-1}$ corresponde ao *reward-to-go* a partir do passo $t$ e $d^{\\pi_\\theta}$ corresponde a distribuição de visitação de estados induzida pela política atual $\\pi_\\theta$. \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class VPG(REINFORCE):\n",
    "    \n",
    "    def __init__(self, obs_space, action_space, config, postprocessing=None):\n",
    "        super(VPG, self).__init__(obs_space, action_space, config, postprocessing)\n",
    "        \n",
    "        config = config[\"value_fn\"]\n",
    "        self.value_fn = build_value_network(obs_space, config[\"hidden_layers\"], config[\"activation\"])\n",
    "        self.value_fn.compile(get_optimizer(config[\"optimizer\"], config[\"learning_rate\"]), loss=\"MSE\")\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",
    "\n",
    "        value_loss = self._train_value_fn(batch)\n",
    "        policy_loss = self._train_policy_fn(batch)\n",
    "      \n",
    "        return policy_loss, value_loss\n",
    "    \n",
    "    def _train_policy_fn(self, batch):\n",
    "        \"\"\"Executa um passo de gradiente ascedente para melhorar a política.\"\"\"\n",
    "        with tf.GradientTape() as tape:\n",
    "            policy_loss = self._loss_pg_fn(batch)\n",
    "            gradients = tape.gradient(policy_loss, self.policy.trainable_weights)\n",
    "        self.optimizer.apply_gradients(zip(gradients, self.policy.trainable_weights))\n",
    "        return policy_loss\n",
    "    \n",
    "    def _train_value_fn(self, batch):\n",
    "        \"\"\"Executa vários passos do gradiente ascedente para melhorar a função Valor (e.g., baseline).\"\"\"\n",
    "        states, rewards = batch[\"states\"], batch[\"rewards\"]\n",
    "        R_t = list(map(self.postprocessing, rewards))\n",
    "        \n",
    "        states = np.concatenate(states, axis=0)\n",
    "        R_t = np.concatenate(R_t, axis=0)\n",
    "\n",
    "        batch_size = len(states)\n",
    "        self.value_fn.fit(states, R_t, epochs=self.config[\"value_fn\"][\"epochs\"], batch_size=batch_size, verbose=0)\n",
    "\n",
    "        value_loss = self.value_fn.evaluate(states, R_t, verbose=0)\n",
    "        return value_loss\n",
    "    \n",
    "    def _loss_pg_fn(self, batch):\n",
    "        \"\"\"Calcula surrogate loss do policy gradient considerando a função de baseline.\"\"\"\n",
    "        states, actions, rewards = batch[\"states\"], batch[\"actions\"], batch[\"rewards\"]\n",
    "        n_episodes = len(states)\n",
    "\n",
    "        loss = 0.0\n",
    "        for episode in range(n_episodes):\n",
    "            action_dist = self.policy(states[episode])\n",
    "            log_prob = action_dist.log_prob(actions[episode])\n",
    "\n",
    "            R_t = self.postprocessing(rewards[episode])\n",
    "            baseline = self.value_fn(states[episode])\n",
    "\n",
    "            loss += - tf.reduce_sum(log_prob * (R_t - baseline))\n",
    "\n",
    "        loss /= n_episodes\n",
    "            \n",
    "        return loss\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Execute o código abaixo rodar o mesmo experimento do **Vanilla Policy Gradients** um número dado de vezes (i.e., `n_trials`):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "n_trials = 5\n",
    "total_timesteps = 1_000_000\n",
    "\n",
    "config = {\n",
    "    **config,\n",
    "    \"value_fn\": {\n",
    "        # network\n",
    "        \"hidden_layers\": [64, 64],\n",
    "        \"activation\": \"relu\",\n",
    "        \n",
    "        # optimization\n",
    "        \"optimizer\": \"rmsprop\",\n",
    "        \"learning_rate\": 3e-3,\n",
    "        \"epochs\": 10,\n",
    "    }\n",
    "}\n",
    "\n",
    "agent_cls = VPG\n",
    "postprocessing = compute_reward_to_go\n",
    "\n",
    "timesteps, total_rewards, avg_total_rewards = utils.runner.run_experiments(\n",
    "    n_trials, env, agent_cls, config, postprocessing, total_timesteps)\n",
    "\n",
    "plot_experiments(env, timesteps, total_rewards, avg_total_rewards)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Execute o código abaixo para visualizar a política aprendida pelo algoritmo **Vanilla Policy Gradient**: "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "n_episodes = 10\n",
    "_ = evaluate(agent, env, n_episodes, render=True)"
   ]
  }
 ],
 "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.6.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}