4a - Generalized Advantage Estimation (GAE) [LunarLander].ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "https://datascience.stackexchange.com/questions/32480/how-does-generalised-advantage-estimation-work\n",
    "https://lilianweng.github.io/lil-log/2018/02/19/a-long-peek-into-reinforcement-learning.html#combining-td-and-mc-learning\n",
    "https://arxiv.org/pdf/1506.02438.pdf\n",
    "https://github.com/higgsfield/RL-Adventure-2\n",
    "http://www.breloff.com/DeepRL-OnlineGAE/\n",
    "https://arxiv.org/pdf/1804.02717.pdf\n",
    "https://ewrl.files.wordpress.com/2015/02/ewrl12_2015_submission_18.pdf\n",
    "https://github.com/Kaixhin/Dist-A3C\n",
    "https://github.com/Kaixhin/Dist-A3C/blob/master/client.py"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "import torch.nn.functional as F\n",
    "import torch.distributions as distributions\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import gym"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_env = gym.make('LunarLander-v2')\n",
    "test_env = gym.make('LunarLander-v2')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "SEED = 1234\n",
    "\n",
    "train_env.seed(SEED);\n",
    "test_env.seed(SEED+1);\n",
    "np.random.seed(SEED);\n",
    "torch.manual_seed(SEED);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "class MLP(nn.Module):\n",
    "    def __init__(self, input_dim, hidden_dim, output_dim, dropout = 0.1):\n",
    "        super().__init__()\n",
    "        \n",
    "        self.net = nn.Sequential(\n",
    "            nn.Linear(input_dim, hidden_dim),\n",
    "            nn.Dropout(dropout),\n",
    "            nn.PReLU(),\n",
    "            nn.Linear(hidden_dim, hidden_dim),\n",
    "            nn.Dropout(dropout),\n",
    "            nn.PReLU(),\n",
    "            nn.Linear(hidden_dim, output_dim)\n",
    "        )\n",
    "        \n",
    "    def forward(self, x):\n",
    "        x = self.net(x)\n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "class ActorCritic(nn.Module):\n",
    "    def __init__(self, actor, critic):\n",
    "        super().__init__()\n",
    "        \n",
    "        self.actor = actor\n",
    "        self.critic = critic\n",
    "        \n",
    "    def forward(self, state):\n",
    "        \n",
    "        action_pred = self.actor(state)\n",
    "        value_pred = self.critic(state)\n",
    "        \n",
    "        return action_pred, value_pred"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "INPUT_DIM = train_env.observation_space.shape[0]\n",
    "HIDDEN_DIM = 128\n",
    "OUTPUT_DIM = train_env.action_space.n\n",
    "\n",
    "actor = MLP(INPUT_DIM, HIDDEN_DIM, OUTPUT_DIM)\n",
    "critic = MLP(INPUT_DIM, HIDDEN_DIM, 1)\n",
    "\n",
    "policy = ActorCritic(actor, critic)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "def init_weights(m):\n",
    "    if type(m) == nn.Linear:\n",
    "        torch.nn.init.xavier_normal_(m.weight)\n",
    "        m.bias.data.fill_(0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "ActorCritic(\n",
       "  (actor): MLP(\n",
       "    (net): Sequential(\n",
       "      (0): Linear(in_features=8, out_features=128, bias=True)\n",
       "      (1): Dropout(p=0.1, inplace=False)\n",
       "      (2): PReLU(num_parameters=1)\n",
       "      (3): Linear(in_features=128, out_features=128, bias=True)\n",
       "      (4): Dropout(p=0.1, inplace=False)\n",
       "      (5): PReLU(num_parameters=1)\n",
       "      (6): Linear(in_features=128, out_features=4, bias=True)\n",
       "    )\n",
       "  )\n",
       "  (critic): MLP(\n",
       "    (net): Sequential(\n",
       "      (0): Linear(in_features=8, out_features=128, bias=True)\n",
       "      (1): Dropout(p=0.1, inplace=False)\n",
       "      (2): PReLU(num_parameters=1)\n",
       "      (3): Linear(in_features=128, out_features=128, bias=True)\n",
       "      (4): Dropout(p=0.1, inplace=False)\n",
       "      (5): PReLU(num_parameters=1)\n",
       "      (6): Linear(in_features=128, out_features=1, bias=True)\n",
       "    )\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "policy.apply(init_weights)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "LEARNING_RATE = 0.0005\n",
    "\n",
    "optimizer = optim.Adam(policy.parameters(), lr = LEARNING_RATE)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train(env, policy, optimizer, discount_factor, trace_decay):\n",
    "    \n",
    "    policy.train()\n",
    "    \n",
    "    log_prob_actions = []\n",
    "    values = []\n",
    "    rewards = []\n",
    "    done = False\n",
    "    episode_reward = 0\n",
    "\n",
    "    state = env.reset()\n",
    "\n",
    "    while not done:\n",
    "\n",
    "        state = torch.FloatTensor(state).unsqueeze(0)\n",
    "\n",
    "        action_pred, value_pred = policy(state)\n",
    "                \n",
    "        action_prob = F.softmax(action_pred, dim = -1)\n",
    "                \n",
    "        dist = distributions.Categorical(action_prob)\n",
    "\n",
    "        action = dist.sample()\n",
    "        \n",
    "        log_prob_action = dist.log_prob(action)\n",
    "        \n",
    "        state, reward, done, _ = env.step(action.item())\n",
    "\n",
    "        log_prob_actions.append(log_prob_action)\n",
    "        values.append(value_pred)\n",
    "        rewards.append(reward)\n",
    "\n",
    "        episode_reward += reward\n",
    "    \n",
    "    log_prob_actions = torch.cat(log_prob_actions)\n",
    "    values = torch.cat(values).squeeze(-1)\n",
    "    \n",
    "    returns = calculate_returns(rewards, discount_factor)\n",
    "    #note: calculate_advantages takes in rewards, not returns!\n",
    "    advantages = calculate_advantages(rewards, values, discount_factor, trace_decay)\n",
    "    \n",
    "    policy_loss, value_loss = update_policy(advantages, log_prob_actions, returns, values, optimizer)\n",
    "\n",
    "    return policy_loss, value_loss, episode_reward"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "def calculate_returns(rewards, discount_factor, normalize = True):\n",
    "    \n",
    "    returns = []\n",
    "    R = 0\n",
    "    \n",
    "    for r in reversed(rewards):\n",
    "        R = r + R * discount_factor\n",
    "        returns.insert(0, R)\n",
    "        \n",
    "    returns = torch.tensor(returns)\n",
    "    \n",
    "    if normalize:\n",
    "        \n",
    "        returns = (returns - returns.mean()) / returns.std()\n",
    "        \n",
    "    return returns"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "def calculate_advantages(rewards, values, discount_factor, trace_decay, normalize = True):\n",
    "    \n",
    "    advantages = []\n",
    "    advantage = 0\n",
    "    next_value = 0\n",
    "    \n",
    "    for r, v in zip(reversed(rewards), reversed(values)):\n",
    "        td_error = r + next_value * discount_factor - v\n",
    "        advantage = td_error + advantage * discount_factor * trace_decay\n",
    "        next_value = v\n",
    "        advantages.insert(0, advantage)\n",
    "        \n",
    "    advantages = torch.tensor(advantages)\n",
    "    \n",
    "    if normalize:\n",
    "        advantages = (advantages - advantages.mean()) / advantages.std()\n",
    "        \n",
    "    return advantages"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "def update_policy(advantages, log_prob_actions, returns, values, optimizer):\n",
    "        \n",
    "    advantages = advantages.detach()\n",
    "    returns = returns.detach()\n",
    "        \n",
    "    policy_loss = - (advantages * log_prob_actions).sum()\n",
    "    \n",
    "    value_loss = F.smooth_l1_loss(returns, values).sum()\n",
    "        \n",
    "    optimizer.zero_grad()\n",
    "    \n",
    "    policy_loss.backward()\n",
    "    value_loss.backward()\n",
    "    \n",
    "    optimizer.step()\n",
    "    \n",
    "    return policy_loss.item(), value_loss.item()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "def evaluate(env, policy):\n",
    "    \n",
    "    policy.eval()\n",
    "    \n",
    "    rewards = []\n",
    "    done = False\n",
    "    episode_reward = 0\n",
    "\n",
    "    state = env.reset()\n",
    "\n",
    "    while not done:\n",
    "\n",
    "        state = torch.FloatTensor(state).unsqueeze(0)\n",
    "\n",
    "        with torch.no_grad():\n",
    "        \n",
    "            action_pred, _ = policy(state)\n",
    "\n",
    "            action_prob = F.softmax(action_pred, dim = -1)\n",
    "                \n",
    "        action = torch.argmax(action_prob, dim = -1)\n",
    "                \n",
    "        state, reward, done, _ = env.step(action.item())\n",
    "\n",
    "        episode_reward += reward\n",
    "        \n",
    "    return episode_reward"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "| Episode:  10 | Mean Train Rewards:  -167.7 | Mean Test Rewards:  -476.3 |\n",
      "| Episode:  20 | Mean Train Rewards:  -209.6 | Mean Test Rewards:  -443.5 |\n",
      "| Episode:  30 | Mean Train Rewards:  -207.5 | Mean Test Rewards:  -276.2 |\n",
      "| Episode:  40 | Mean Train Rewards:  -184.9 | Mean Test Rewards:     9.2 |\n",
      "| Episode:  50 | Mean Train Rewards:  -147.1 | Mean Test Rewards:   -20.4 |\n",
      "| Episode:  60 | Mean Train Rewards:  -123.5 | Mean Test Rewards:   -96.6 |\n",
      "| Episode:  70 | Mean Train Rewards:  -128.2 | Mean Test Rewards:   -45.4 |\n",
      "| Episode:  80 | Mean Train Rewards:  -128.0 | Mean Test Rewards:   -30.5 |\n",
      "| Episode:  90 | Mean Train Rewards:  -128.4 | Mean Test Rewards:   -75.3 |\n",
      "| Episode: 100 | Mean Train Rewards:  -110.3 | Mean Test Rewards:  -208.3 |\n",
      "| Episode: 110 | Mean Train Rewards:  -115.3 | Mean Test Rewards:  -362.8 |\n",
      "| Episode: 120 | Mean Train Rewards:   -96.6 | Mean Test Rewards:  -607.0 |\n",
      "| Episode: 130 | Mean Train Rewards:  -100.1 | Mean Test Rewards:  -985.8 |\n",
      "| Episode: 140 | Mean Train Rewards:   -98.6 | Mean Test Rewards: -1331.9 |\n",
      "| Episode: 150 | Mean Train Rewards:  -117.8 | Mean Test Rewards: -1366.6 |\n",
      "| Episode: 160 | Mean Train Rewards:  -120.6 | Mean Test Rewards: -1356.4 |\n",
      "| Episode: 170 | Mean Train Rewards:  -118.9 | Mean Test Rewards: -1389.0 |\n",
      "| Episode: 180 | Mean Train Rewards:  -107.4 | Mean Test Rewards: -1557.8 |\n",
      "| Episode: 190 | Mean Train Rewards:  -109.4 | Mean Test Rewards: -1638.7 |\n",
      "| Episode: 200 | Mean Train Rewards:  -115.1 | Mean Test Rewards: -1551.3 |\n",
      "| Episode: 210 | Mean Train Rewards:   -96.5 | Mean Test Rewards: -1780.7 |\n",
      "| Episode: 220 | Mean Train Rewards:   -70.0 | Mean Test Rewards: -2215.9 |\n",
      "| Episode: 230 | Mean Train Rewards:   -75.8 | Mean Test Rewards: -2133.9 |\n",
      "| Episode: 240 | Mean Train Rewards:   -71.7 | Mean Test Rewards: -2612.4 |\n",
      "| Episode: 250 | Mean Train Rewards:   -90.6 | Mean Test Rewards: -2863.6 |\n",
      "| Episode: 260 | Mean Train Rewards:   -68.9 | Mean Test Rewards: -4042.3 |\n",
      "| Episode: 270 | Mean Train Rewards:   -76.4 | Mean Test Rewards: -3659.9 |\n",
      "| Episode: 280 | Mean Train Rewards:   -71.2 | Mean Test Rewards: -2684.7 |\n",
      "| Episode: 290 | Mean Train Rewards:   -85.1 | Mean Test Rewards: -1661.0 |\n",
      "| Episode: 300 | Mean Train Rewards:   -78.1 | Mean Test Rewards: -1253.0 |\n",
      "| Episode: 310 | Mean Train Rewards:   -76.8 | Mean Test Rewards: -1201.1 |\n",
      "| Episode: 320 | Mean Train Rewards:   -57.7 | Mean Test Rewards: -1318.0 |\n",
      "| Episode: 330 | Mean Train Rewards:   -50.3 | Mean Test Rewards: -1219.0 |\n",
      "| Episode: 340 | Mean Train Rewards:   -41.7 | Mean Test Rewards: -1110.5 |\n",
      "| Episode: 350 | Mean Train Rewards:   -33.0 | Mean Test Rewards:  -811.2 |\n",
      "| Episode: 360 | Mean Train Rewards:   -42.7 | Mean Test Rewards:  -877.1 |\n",
      "| Episode: 370 | Mean Train Rewards:   -72.8 | Mean Test Rewards: -1050.2 |\n",
      "| Episode: 380 | Mean Train Rewards:   -89.9 | Mean Test Rewards: -1137.1 |\n",
      "| Episode: 390 | Mean Train Rewards:   -78.1 | Mean Test Rewards: -1400.2 |\n",
      "| Episode: 400 | Mean Train Rewards:   -67.3 | Mean Test Rewards: -1688.1 |\n",
      "| Episode: 410 | Mean Train Rewards:   -58.8 | Mean Test Rewards: -1565.0 |\n",
      "| Episode: 420 | Mean Train Rewards:   -50.1 | Mean Test Rewards: -1014.3 |\n",
      "| Episode: 430 | Mean Train Rewards:   -34.2 | Mean Test Rewards:  -797.6 |\n",
      "| Episode: 440 | Mean Train Rewards:   -19.9 | Mean Test Rewards:  -511.5 |\n",
      "| Episode: 450 | Mean Train Rewards:    -1.5 | Mean Test Rewards:  -274.3 |\n",
      "| Episode: 460 | Mean Train Rewards:    -5.0 | Mean Test Rewards:  -201.7 |\n",
      "| Episode: 470 | Mean Train Rewards:    10.6 | Mean Test Rewards:  -265.2 |\n",
      "| Episode: 480 | Mean Train Rewards:     1.5 | Mean Test Rewards:  -328.7 |\n",
      "| Episode: 490 | Mean Train Rewards:    -5.8 | Mean Test Rewards:  -342.8 |\n",
      "| Episode: 500 | Mean Train Rewards:   -15.7 | Mean Test Rewards:  -385.6 |\n",
      "| Episode: 510 | Mean Train Rewards:   -13.3 | Mean Test Rewards:  -572.9 |\n",
      "| Episode: 520 | Mean Train Rewards:    10.7 | Mean Test Rewards:  -766.1 |\n",
      "| Episode: 530 | Mean Train Rewards:     9.8 | Mean Test Rewards:  -641.0 |\n",
      "| Episode: 540 | Mean Train Rewards:     7.0 | Mean Test Rewards:  -344.2 |\n",
      "| Episode: 550 | Mean Train Rewards:    -3.2 | Mean Test Rewards:  -240.4 |\n"
     ]
    }
   ],
   "source": [
    "MAX_EPISODES = 1_000\n",
    "DISCOUNT_FACTOR = 0.99\n",
    "TRACE_DECAY = 0.99\n",
    "N_TRIALS = 25\n",
    "REWARD_THRESHOLD = 200\n",
    "PRINT_EVERY = 10\n",
    "\n",
    "train_rewards = []\n",
    "test_rewards = []\n",
    "\n",
    "for episode in range(1, MAX_EPISODES+1):\n",
    "    \n",
    "    policy_loss, value_loss, train_reward = train(train_env, policy, optimizer, DISCOUNT_FACTOR, TRACE_DECAY)\n",
    "    \n",
    "    test_reward = evaluate(test_env, policy)\n",
    "    \n",
    "    train_rewards.append(train_reward)\n",
    "    test_rewards.append(test_reward)\n",
    "    \n",
    "    mean_train_rewards = np.mean(train_rewards[-N_TRIALS:])\n",
    "    mean_test_rewards = np.mean(test_rewards[-N_TRIALS:])\n",
    "    \n",
    "    if episode % PRINT_EVERY == 0:\n",
    "    \n",
    "        print(f'| Episode: {episode:3} | Mean Train Rewards: {mean_train_rewards:7.1f} | Mean Test Rewards: {mean_test_rewards:7.1f} |')\n",
    "    \n",
    "    if mean_test_rewards >= REWARD_THRESHOLD:\n",
    "        \n",
    "        print(f'Reached reward threshold in {episode} episodes')\n",
    "        \n",
    "        break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure(figsize=(12,8))\n",
    "plt.plot(test_rewards, label='Test Reward')\n",
    "plt.plot(train_rewards, label='Train Reward')\n",
    "plt.xlabel('Episode', fontsize=20)\n",
    "plt.ylabel('Reward', fontsize=20)\n",
    "plt.hlines(REWARD_THRESHOLD, 0, len(test_rewards), color='r')\n",
    "plt.legend(loc='lower right')\n",
    "plt.grid()"
   ]
  }
 ],
 "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": 2
}