04_PG.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Policy Gradient Methods\n",
    "\n",
    "Policy _based_ methods learn the optimal policy directly, without necessarily estimating a value\n",
    "function. Policy _gradient_ methods do that performing gradient ascent on the objective function.\n",
    "\n",
    "### Advantages\n",
    "\n",
    " * No need to store action-values.\n",
    " * Ability to learn a stochastic policy direcly.\n",
    " * Hence, no need to manually tune exploitation vs. exploration.\n",
    " * Effective in continuous action spaces (and high-dimensional state spaces).\n",
    " * They generally have good convergence properties.\n",
    "\n",
    "### Disadvantages\n",
    "\n",
    " * They might have high-variance.\n",
    " * Might converge to a local maximum.\n",
    " * Slower than other methods, and might take a long time to train."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import torch\n",
    "import torch.optim as optim\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "from torch.distributions import Categorical\n",
    "from typing import Union\n",
    "\n",
    "import gymnasium as gym\n",
    "\n",
    "from util.gymnastics import DEVICE, gym_simulation, init_random, plot_scores"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## CartPole Environment"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# `gym_simulation` is a convenient utility to run simulations on Gynmasium environment, so that we\n",
    "# don't have to repeat the same code in all notebooks. But it works basically the same as the one we\n",
    "# implemented for DQN :) Feel free to check the source code!\n",
    "\n",
    "gym_simulation(\"CartPole-v1\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Just for convenience, we hardcode the state and action sizes of the CartPole environment.\n",
    "STATE_SIZE  = 4\n",
    "ACTION_SIZE = 2"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Optimization Rule\n",
    "\n",
    "For one trajectory $\\tau$ (or episode), the neural networks weight can be updated according to:\n",
    "\n",
    "$$\n",
    "\\theta_{k+1} = \\theta_k - \\alpha \\sum_{t=0} \\nabla_{\\theta} \\log \\pi_{\\theta}(a_t|s_t) R(\\tau)\n",
    "$$\n",
    "\n",
    "We can interpret this as pushing up probabilities for action / states combinations when the return\n",
    "is high, and the other way around for low returns.\n",
    "\n",
    "This relationship is also interesting because only the policy function needs to be differentiable:\n",
    "the reward function might very well be discontinuous and sparse.\n",
    "\n",
    "For derivation, check the [Hugging Face Deep RL tutorial](https://huggingface.co/learn/deep-rl-course/unit4/pg-theorem)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## REINFORCE\n",
    "\n",
    "<div style=\"width: 70%\">\n",
    "  <img src=\"assets/04_PG_reinforce.png\">\n",
    "  <br>\n",
    "  <small>Sutton & Barto 2022</small>\n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "class PolicyNetwork(nn.Module):\n",
    "    \"\"\"The neural network computing the stochastic policy (action probabilities for a state).\"\"\"\n",
    "    def __init__(self, hidden_units=16):\n",
    "        super(PolicyNetwork, self).__init__()\n",
    "        # TODO: Create two fully connected / linear layers, input dimension STATE_SIZE, output\n",
    "        #       dimension ACTION_SIZE, and hidden units specified in the constructor.\n",
    "\n",
    "    def forward(self, x):\n",
    "        # TODO: Use ReLU as first non-linearity, and softmax for the output layer (so that we can\n",
    "        #       interpret this as a stochastic policy across action probabilities). Hint: make sure\n",
    "        #       to choose the right dimension for the softmax!\n",
    "        pass"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Agent:\n",
    "    def __init__(self):\n",
    "        self.policy = PolicyNetwork()\n",
    "        self.optimizer = optim.Adam(self.policy.parameters(), lr=1e-2)\n",
    "\n",
    "    def sample_action(self, state: np.array):\n",
    "        \"\"\"Samples an action for a state from the policy network.\"\"\"\n",
    "        # TODO: Convert the state to a PyTorch tensor.\n",
    "        state = None\n",
    "        # TODO: Get the action probabilities from the policy.\n",
    "        probs = None\n",
    "        # TODO: Create a Categorical distribution with those probabilities.\n",
    "        cdist = None\n",
    "        # TODO: Sample the action from the categorical distribution.\n",
    "        action = None\n",
    "        # TODO: Return the action (hint: use item()), and the log_prob for that action (hint: use\n",
    "        #       the distribution again!)\n",
    "        return None\n",
    "\n",
    "    def learn(self, log_probs: list[torch.Tensor], returns: Union[np.float64, np.array]):\n",
    "        \"\"\"Perform a step of learning (gradient ascent) on the policy network.\"\"\"\n",
    "        # TODO: For reasons you'll see below, rewards can be either a scalar or an array. Let's\n",
    "        #       create a corresponding tensor first.\n",
    "        returns = None\n",
    "\n",
    "        # TODO: Compute the policy loss. Hint: gradient ascent.... so negative sign!\n",
    "        policy_loss = None\n",
    "\n",
    "        # TODO: Perform a backprop step via the optimizer.\n",
    "    \n",
    "    @torch.no_grad\n",
    "    def act(self, state):\n",
    "        \"\"\"Convenient method for the agent to select an action during simulation.\"\"\"\n",
    "        return self.sample_action(state)[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "def REINFORCE(env, agent, max_episodes=10_000, max_t=1_000, gamma=1.0):\n",
    "    # Tracks the score for each episode.\n",
    "    scores = []\n",
    "    # Loop for max_episodes.\n",
    "    for i_episode in range(1, max_episodes + 1):\n",
    "        # Store episode rewards.\n",
    "        rewards = []\n",
    "        # Store episode log probabilities\n",
    "        log_probs = []\n",
    "        # Start the episode in the initial state.\n",
    "        state, _ = env.reset()\n",
    "\n",
    "        # TODO: Generate an episode following the policy, for T (max_t) timesteps.\n",
    "        for _ in range(max_t):\n",
    "            # TODO: sample action and log probability.\n",
    "            action, log_prob = None\n",
    "            # TODO: perform a step in the environment.\n",
    "            state, reward, terminated, truncated, _ = None\n",
    "            # TODO: store reward and log probability.\n",
    "            # TODO: Check for episode completion.\n",
    "\n",
    "        # Compute discounted return.\n",
    "        # TODO: First, compute the discounts.\n",
    "        discounts = None\n",
    "        # TODO: Then compute the total discounted return as the sum of the discouted rewards.\n",
    "        R = None\n",
    "\n",
    "        # TODO: Perform a learning step on the agent calling the `learn` method.\n",
    "\n",
    "        # Track scores and print statistics.\n",
    "        scores.append(sum(rewards))\n",
    "        avg_score = np.mean(scores[-100:])\n",
    "        if i_episode % 100 == 0:\n",
    "            print(f'Episode {i_episode}\\tAverage Score: {avg_score:.2f}')\n",
    "        if avg_score >= 490.0: # Solved\n",
    "            print(f'Environment solved at episode {i_episode}\\tAverage Score: {avg_score:.2f}')\n",
    "            break\n",
    "\n",
    "    return scores"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "with init_random(gym.make('CartPole-v1')) as env:\n",
    "    agent = Agent()\n",
    "    scores = REINFORCE(env, agent)\n",
    "plot_scores(scores)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "gym_simulation(\"CartPole-v1\", agent)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Improvements\n",
    "\n",
    "### Use Future Rewards\n",
    "\n",
    "First thing to notice is that we are using all rewards at every timestep. But really, we should only\n",
    "consider future rewards: i.e., the rewards that are actually the consequences of our actions.\n",
    "\n",
    "$$\n",
    "g = \\sum_t R_t^{future}\\nabla_{\\theta} log\\pi_{\\theta}(a_t | s_t)\n",
    "$$\n",
    "\n",
    "### Normalize Rewards\n",
    "\n",
    "Another technique we can use (especially when collecting multiple trajectories, that will come later\n",
    "on!) is to normalize the rewards: which roughly picks half actions to encourage / discourage, and\n",
    "keeps the gradient updates moderate.\n",
    "\n",
    "$$\n",
    "R_k \\leftarrow \\frac{R_k - \\mu}{\\sigma}\n",
    "$$\n",
    "\n",
    "### Baseline Subtraction\n",
    "\n",
    "The idea is to subtract to the reward a _baseline_ $b$, for example the average reward along all\n",
    "trajectories (What if every trajectory has _always_ positive returns?). In this case, things that\n",
    "are above average will push their probabilities to happen up while things below average will be\n",
    "penalized.\n",
    "\n",
    "$$\n",
    "\\theta_{k+1} = \\theta_k - \\alpha \\sum_{t=0} \\nabla_{\\theta} \\log \\pi_{\\theta}(a_t|s_t) [R(\\tau) - b]\n",
    "$$\n",
    "\n",
    "We can do this because on expectation this extra subtracted term will have zero effect, but overall\n",
    "we'll get reduced variance (proof left as exercise and/or you can find it in the resources).\n",
    "\n",
    "#### Advantage Function\n",
    "\n",
    "This value that we multiply with the log-probability to \"reinforce\" or \"depress\" the corresponding\n",
    "actions is called the _advantage function_ and plays a critical role in state-of-the-art algorithms:\n",
    "\n",
    "$$\n",
    "A(*) = R(\\tau) - b\n",
    "$$\n",
    "\n",
    "It measures how better the selected action does compared to the _average_ in the state."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "def REINFORCE_v2(env, agent, max_episodes=10_000, max_t=1_000, gamma=1.0):\n",
    "    # Tracks the score for each episode.\n",
    "    scores = []\n",
    "    # Loop for max_episodes.\n",
    "    for i_episode in range(1, max_episodes + 1):\n",
    "        # Store episode rewards.\n",
    "        rewards = []\n",
    "        # Store episode log probabilities\n",
    "        log_probs = []\n",
    "        # Start the episode in the initial state.\n",
    "        state, _ = env.reset()\n",
    "\n",
    "        # TODO: Generate an episode following the policy. Copy the code above :)\n",
    "        for _ in range(max_t):\n",
    "            pass\n",
    "\n",
    "        # Compute discounted _future_ returns.\n",
    "        # TODO: Compute the discounts and discounted rewards.\n",
    "        discounted_rewards = None\n",
    "        # TODO: Compute the future_returns. Hint: consider cumulative sums in reverse order :)\n",
    "        future_returns = None\n",
    "\n",
    "        # Baseline.\n",
    "        # TODO: Use the average of future returns as baseline.\n",
    "        baseline = None\n",
    "        # TODO: Subtract the baseline from the future returns.\n",
    "        future_returns = None\n",
    "\n",
    "        # Normalization.\n",
    "        # TODO: normalize the returns computing mean and standard deviation. Hint: make sure the std\n",
    "        #       is non zero; use np.mean and np.std.\n",
    "        normalized_returns = None\n",
    "\n",
    "        # TODO: Perform a learning step calling agent.learn(...)\n",
    "        # copy() for negative strides :(\n",
    "        #   https://discuss.pytorch.org/t/negative-strides-in-tensor-error/134287/2\n",
    "\n",
    "        # Track scores and print statistics\n",
    "        scores.append(sum(rewards))\n",
    "        avg_score = np.mean(scores[-100:])\n",
    "        if i_episode % 100 == 0:\n",
    "            print(f'Episode {i_episode}\\tAverage Score: {avg_score:.2f}')\n",
    "        if avg_score >= 490.0: # Solved\n",
    "            print(f'Environment solved at episode {i_episode}\\tAverage Score: {avg_score:.2f}')\n",
    "            break\n",
    "\n",
    "    return scores"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with init_random(gym.make('CartPole-v1')) as env:\n",
    "    agent_v2 = Agent()\n",
    "    scores_v2 = REINFORCE_v2(env, agent_v2)\n",
    "plot_scores(scores_v2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "gym_simulation(\"CartPole-v1\", agent_v2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## What Can We Do Better?\n",
    "\n",
    "There are other improvements that can be applied:\n",
    "\n",
    "* To _reduce noise_ on the gradient, we can simply sample multiple different trajectories and learn\n",
    "  from all of those. Vectorized environments will help with this!\n",
    "* Actor-critic setup and advanced advantage estimation such as _GAE_ will improve learning.\n",
    "* We are currently _discarding experiences_ after every learning step. That is because the policy\n",
    "  effectively changes. But we'll see that with importance sampling we can iterate on the same data\n",
    "  multiple times and learn in mini-batches!\n",
    "* Techniques such as \"trust region\" and \"_loss clipping_\" will help against degeneration and keep\n",
    "  the policy learning along smooth gradient directions.\n",
    "\n",
    "Once we put all of these in place... we'll have PPO!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Appendix\n",
    "\n",
    "### Meaning of Loss\n",
    "\n",
    "Note that the loss function used in policy gradient methods doesn't have the same meaning of the\n",
    "typical supervised learning setup. In particular, after that first step of gradient descent, there\n",
    "is no more connection to performance - which is determined by the average return.\n",
    "\n",
    "The loss function is only useful when evaluated at the current parameters to perform one step of\n",
    "gradient ascent. After that it loses its meaning and it's value shouldn't be used as a metric for\n",
    "performance.\n",
    "\n",
    "More details on [OpenAI Spinning Up](https://spinningup.openai.com/en/latest/spinningup/rl_intro3.html#id14).\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "drlzh",
   "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.11.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}