2_ActorCritic.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "2_ActorCritic.ipynb",
      "provenance": [],
      "collapsed_sections": [],
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/tzs930/mlbootcamp-2022-rl-practice/blob/main/2_ActorCritic.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Reinforcement Learning Practice 2 : REINFORCE, Actor-Critic\n",
        "\n",
        "- In this assignment, we will implement two basic policy gradient methods, REINFORCE and (state-value) actor-critic."
      ],
      "metadata": {
        "id": "HJ8P1GmOXL1O"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "!pip install numpy torch matplotlib gym"
      ],
      "metadata": {
        "id": "89AyTVW2wm6r"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# import packages\n",
        "import torch\n",
        "import torch.nn as nn\n",
        "import torch.nn.functional as F\n",
        "import torch.optim as optim\n",
        "from torch.distributions import Categorical\n",
        "import numpy as np    \n",
        "import gym\n",
        "import matplotlib.pyplot as plt\n",
        "\n",
        "np.random.seed(123)\n",
        "torch.manual_seed(123)"
      ],
      "metadata": {
        "id": "I8UdOQN3yipP"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "This assignment features the Cartpole domain which tasks the agent with balancing a pole affixed to a movable cart. The agent employs two discrete actions which apply force to the cart. Episodes provide +1 reward for each step in which the pole has not fallen over, up to a maximum of 200 steps. (See https://gym.openai.com/envs/CartPole-v0/) for more details. \n",
        "\n",
        "Additionally, we'll save the size of the state and action spaces, and define hyperparameters such as the number of hidden units in our network. These parameters don't need to be changed, but you can try varying hyperparameters and see how learning is affected."
      ],
      "metadata": {
        "id": "g1b8nW6vfc5S"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "env = gym.make('CartPole-v0')\n",
        "env.seed(123)\n",
        "\n",
        "state_dim = env.observation_space.shape[0] # Dimension of state space\n",
        "action_count = env.action_space.n          # Number of actions\n",
        "\n",
        "# Hyperparameters\n",
        "hidden_size = 128             # Number of hidden units\n",
        "max_number_of_episodes = 500  # Number of training episodes\n",
        "log_frequency = 20            # Frequency of logging\n",
        "gamma = 0.999                 # Discount rate\n",
        "policy_lr = 1e-2              # Learning rate of policy (actor) network\n",
        "critic_lr = 1e-2              # Learning rate of critic network"
      ],
      "metadata": {
        "id": "xtdysW6VXTxj"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## 1. Implementing REINFORCE algorithm\n",
        "- In this section, we will implement REINFORCE algorithm.\n",
        "- We use the softmax policy since Cartpole domain has the discrete action space.\n",
        "- Recall that REINFORCE objective is: \n",
        "$$\\sum_{t=1}^T \\nabla_\\theta \\log \\pi_\\theta(a_t|s_t) R_t$$ where $R_t = \\sum_{i=t}^{T} \\gamma^i r_i$.\n",
        "- Then, the desired loss function for REINFORCE is: \n",
        "$$L_{\\pi}(\\theta) = \\sum_{t=1}^T \\log \\pi_\\theta(a_t|s_t) R_t$$.\n"
      ],
      "metadata": {
        "id": "05YsGQ4VfMeY"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "class Policy(nn.Module):\n",
        "    def __init__(self, state_dim, action_count, hidden_size):\n",
        "        super(Policy, self).__init__()\n",
        "        self.W1 = nn.Linear(state_dim, hidden_size)\n",
        "        self.W2 = nn.Linear(hidden_size, action_count)\n",
        "\n",
        "    def forward(self, state):\n",
        "        out = self.W1(state)\n",
        "        out = F.relu(out)\n",
        "        out = self.W2(out)\n",
        "        return F.softmax(out, dim=-1)"
      ],
      "metadata": {
        "id": "Hw7CODARXU58"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Now, we have two things to do here:\n",
        "- TODO 1 : define **the list of discounted returns** for the loss function \n",
        "\n",
        " (Hint: we need $[R_1, R_2, ..., R_T]$ where $R_t = \\sum_{i=t}^{T} \\gamma^i r_i$)\n",
        "- TODO 2 : define the loss for **REINFORCE** algorithm\n",
        "\n",
        " (Hint: we can easily obtain loss function using the inner product of discounted returns and log-probabilities, i.e.  $[\\log \\pi(a_1|s_1), \\log \\pi(a_2|s_2), ..., \\log \\pi(a_T|s_T)]^\\top [R_1, R_2, ..., R_T]= \\sum_{t=1}^T \\log \\pi(a_t|s_t) R_t$ )"
      ],
      "metadata": {
        "id": "564arZqDjfCO"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "policy = Policy(state_dim, action_count, hidden_size)\n",
        "policy_optimizer = optim.Adam(policy.parameters(), lr=policy_lr)\n",
        "\n",
        "reward_sum = 0\n",
        "episode_return_list = []\n",
        "episode_length_list = []\n",
        "\n",
        "num_episode_list = []\n",
        "reinforce_score_avg = []\n",
        "reinforce_score_std = []\n",
        "\n",
        "states, rewards, actions, logprobs = [],[],[],[]\n",
        "    \n",
        "for episode_number in range(max_number_of_episodes):\n",
        "    episode_return = 0\n",
        "    episode_length = 0\n",
        "    done = False\n",
        "    observation = env.reset()\n",
        "    t = 1\n",
        "    while not done:\n",
        "        state = np.reshape(observation, [1, state_dim]).astype(np.float32)\n",
        "        states.append(state)\n",
        "\n",
        "        # Run the policy network and get an action to take.\n",
        "        state = torch.Tensor(state)\n",
        "        probs = policy(state)[0]\n",
        "        dist = Categorical(probs)\n",
        "        action = dist.sample()\n",
        "        logprob = dist.log_prob(action)\n",
        "        action = action.detach().numpy()\n",
        "\n",
        "        logprobs.append(logprob)\n",
        "        \n",
        "        # step the environment and get new measurements\n",
        "        observation, reward, done, _ = env.step(action)\n",
        "        reward_sum += float(reward)\n",
        "\n",
        "        # Record reward (has to be done after we call step() to get reward for previous action)\n",
        "        rewards.append(float(reward))\n",
        "        \n",
        "        episode_return += reward\n",
        "        episode_length = t\n",
        "        t += 1\n",
        "\n",
        "    episode_return_list.append(episode_return)\n",
        "    episode_length_list.append(episode_length)\n",
        "\n",
        "    # Finish Episode\n",
        "    # Compute the discounted reward backwards through time.\n",
        "    R = 0\n",
        "    returns = []\n",
        "    ##########################################################################################\n",
        "    # TODO 1:\n",
        "    # - define a list of discounted sums\n",
        "    # i.e. [r_0 + \\gamma * r_1 + ... + \\gamma^{T-1} * r_T, ... ,  r_{T-1} + \\gamma * r_T, r_T]\n",
        "    for r in rewards[::-1]:\n",
        "      R = r\n",
        "      returns.insert(0, R)\n",
        "    ##########################################################################################\n",
        "\n",
        "    returns = torch.tensor(returns)\n",
        "    \n",
        "    policy_loss = 0\n",
        "    for log_prob, R in zip(logprobs, returns):\n",
        "      ########################################################################################\n",
        "      # TODO 2:\n",
        "      # - define REINFORCE loss\n",
        "      policy_loss += logprob\n",
        "      ########################################################################################\n",
        "    \n",
        "    policy_optimizer.zero_grad()\n",
        "    policy_loss.backward()\n",
        "    policy_optimizer.step()\n",
        "    \n",
        "    states, rewards, actions, logprobs = [],[],[],[]\n",
        "\n",
        "    if episode_number % log_frequency == 0:\n",
        "        print('Episode: %d. Average reward for episode %f. Variance %f' % (episode_number, np.mean(episode_return_list), np.std(episode_return_list)**2 ))\n",
        "        num_episode_list.append(episode_number)\n",
        "        reinforce_score_avg.append(np.mean(episode_return_list))\n",
        "        reinforce_score_std.append(np.std(episode_return_list))\n",
        "        episode_return_list = []\n"
      ],
      "metadata": {
        "id": "ffmZ6Ux3XWnc"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "- After training, plot the training curve using the below code."
      ],
      "metadata": {
        "id": "mxneu7Gi4Gmo"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "num_episode_list = np.arange(0,max_number_of_episodes,log_frequency)\n",
        "plt.plot(num_episode_list, reinforce_score_avg, label='REINFORCE')\n",
        "plt.fill_between(x=num_episode_list,\n",
        "                 y1=np.array(reinforce_score_avg)-np.array(reinforce_score_std),\n",
        "                 y2=np.array(reinforce_score_avg)+np.array(reinforce_score_std),\n",
        "                 alpha=0.3)\n",
        "plt.legend()\n",
        "plt.xlabel('The number of episodes')\n",
        "plt.ylabel('Episode score')"
      ],
      "metadata": {
        "id": "ifl7iI774C45"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## 2. Implementing Actor-Critic algorithm\n",
        "- In this section, we will implement Actor-Critic algorithm, especially using the state-value function as a baseline for reducing variance.\n",
        "- We will approximate the state-value function $V_\\phi(s)$ with the discounted return of the remained epsiode starting from $s$, i.e., \n",
        "$$\n",
        "V_\\phi(s) \\approx V^\\pi(s) = \\mathbb{E}_\\pi\\left[\\sum_{i=t}^T \\gamma^i r_i | s_t=s \\right]\n",
        "$$\n",
        "- Recall that \n",
        " - the desired loss function for the critic is L2 loss over the episode return:\n",
        " $L_{V}(\\phi) = \\frac{1}{m}\\sum_{t=1}^m  (V_\\phi(s_t) - R_t)^2$.\n",
        "\n",
        " - the desired loss function for the actor is: \n",
        " $L_\\pi(\\theta) = \\frac{1}{m}\\sum_{t=1}^m \\log \\pi_\\theta(a_t|s_t) (R_t - V_\\phi(s_t))$.\n",
        " \n"
      ],
      "metadata": {
        "id": "53fvUlOxgyWT"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "class Critic(nn.Module):\n",
        "    def __init__(self, state_dim, hidden_size):\n",
        "        super(Critic, self).__init__()\n",
        "        self.W1 = nn.Linear(state_dim, hidden_size)\n",
        "        self.W2 = nn.Linear(hidden_size, 1)\n",
        "\n",
        "    def forward(self, state) :\n",
        "        out = self.W1(state)\n",
        "        out = F.relu(out)\n",
        "        out = self.W2(out)\n",
        "        \n",
        "        return out"
      ],
      "metadata": {
        "id": "vzKW4w6zevmV"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "We have two things to do here:\n",
        "- TODO 1: define the list of discounted returns (reuse the result of TODO 1)\n",
        "- TODO 3: define **the critic loss** and **the actor loss**."
      ],
      "metadata": {
        "id": "W87GUWqt4djV"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "actor = Policy(state_dim, action_count, hidden_size)\n",
        "actor_optimizer = optim.Adam(actor.parameters(), lr=policy_lr)\n",
        "critic = Critic(state_dim, hidden_size)\n",
        "critic_optimizer = optim.Adam(critic.parameters(), lr=critic_lr)\n",
        "\n",
        "reward_sum = 0\n",
        "\n",
        "max_number_of_episodes = 500\n",
        "episode_return_list = []\n",
        "episode_length_list = []\n",
        "\n",
        "num_episode_list = []\n",
        "ac_score_avg = []\n",
        "ac_score_std = []\n",
        "\n",
        "states, rewards, actions, logprobs, state_values = [],[],[],[],[]\n",
        "    \n",
        "for episode_number in range(max_number_of_episodes):\n",
        "    episode_return = 0\n",
        "    episode_length = 0\n",
        "    done = False\n",
        "    observation = env.reset()\n",
        "    t = 1\n",
        "    while not done:\n",
        "        state = np.reshape(observation, [1, state_dim]).astype(np.float32)\n",
        "        states.append(state)\n",
        "\n",
        "        # Run the policy network and get an action to take.\n",
        "        state = torch.Tensor(state)\n",
        "        probs = actor(state)[0]\n",
        "        dist = Categorical(probs)\n",
        "        action = dist.sample()\n",
        "        logprob = dist.log_prob(action)\n",
        "        action = action.detach().numpy()\n",
        "        state_value = critic(state)\n",
        "\n",
        "        logprobs.append(logprob)\n",
        "        state_values.append(state_value)\n",
        "        \n",
        "        # step the environment and get new measurements\n",
        "        observation, reward, done, _ = env.step(action)\n",
        "        reward_sum += float(reward)\n",
        "\n",
        "        # Record reward (has to be done after we call step() to get reward for previous action)\n",
        "        rewards.append(float(reward))\n",
        "        \n",
        "        episode_return += reward\n",
        "        episode_length = t\n",
        "        t += 1\n",
        "\n",
        "    episode_return_list.append(episode_return)\n",
        "    episode_length_list.append(episode_length)\n",
        "\n",
        "    # Finish Episode\n",
        "    # Compute the discounted reward backwards through time.\n",
        "    R = 0\n",
        "    returns = []\n",
        "    ##########################################################################################\n",
        "    # TODO 1: Reuse result of TODO 1 for discounted return\n",
        "    for r in rewards[::-1]:\n",
        "      R = r\n",
        "      returns.insert(0, R)\n",
        "    ##########################################################################################\n",
        "    \n",
        "    returns = torch.tensor(returns)\n",
        "    actor_loss = 0\n",
        "    critic_loss = 0\n",
        "\n",
        "    for log_prob, state_value, R in zip(logprobs, state_values, returns):        \n",
        "        ##########################################################################################\n",
        "        # TODO 3:\n",
        "        # - define actor_loss and critic_loss here\n",
        "        actor_loss += log_prob\n",
        "        critic_loss += state_value\n",
        "        ##########################################################################################\n",
        "\n",
        "    actor_optimizer.zero_grad()\n",
        "    actor_loss.backward()\n",
        "    actor_optimizer.step()\n",
        "\n",
        "    critic_optimizer.zero_grad()\n",
        "    critic_loss.backward()\n",
        "    critic_optimizer.step()\n",
        "    \n",
        "    states, rewards, actions, logprobs, state_values = [],[],[],[],[]\n",
        "\n",
        "    if episode_number % log_frequency == 0:\n",
        "        print('Episode: %d. Average reward for episode %f. Variance %f' % (episode_number, np.mean(episode_return_list), np.std(episode_return_list)**2 ))\n",
        "        ac_score_avg.append(np.mean(episode_return_list))\n",
        "        ac_score_std.append(np.std(episode_return_list))\n",
        "        episode_return_list = []"
      ],
      "metadata": {
        "id": "2agAZTdTXiJN"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "- After training, plot and compare training curves of REINFORCE and Actor-Critic algorithm using the below code."
      ],
      "metadata": {
        "id": "CsE2EuwL57rn"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "num_episode_list = np.arange(0,max_number_of_episodes,log_frequency)\n",
        "plt.plot(num_episode_list, reinforce_score_avg, label='REINFORCE')\n",
        "plt.fill_between(x=num_episode_list,\n",
        "                 y1=np.array(reinforce_score_avg)-np.array(reinforce_score_std),\n",
        "                 y2=np.array(reinforce_score_avg)+np.array(reinforce_score_std),\n",
        "                 alpha=0.3)\n",
        "plt.plot(num_episode_list, ac_score_avg, label='Actor-Critic')\n",
        "plt.fill_between(x=num_episode_list,\n",
        "                 y1=np.array(ac_score_avg)-np.array(ac_score_std),\n",
        "                 y2=np.array(ac_score_avg)+np.array(ac_score_std),\n",
        "                 alpha=0.3)\n",
        "plt.legend()\n",
        "plt.xlabel('The number of episodes')\n",
        "plt.ylabel('Episode score')"
      ],
      "metadata": {
        "id": "x_llCo68XuqQ"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}