exercises/templates/ex12/DDPG_and_PPO.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "88d5a01d",
   "metadata": {
    "nbgrader": {
     "grade": false,
     "grade_id": "cell-0ce69fe260e748d8",
     "locked": true,
     "schema_version": 3,
     "solution": false,
     "task": false
    }
   },
   "source": [
    "# Exercice 12) Deep Deterministic Policy Gradients and Proximal Policy Optimization"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ea562d4a",
   "metadata": {
    "nbgrader": {
     "grade": false,
     "grade_id": "cell-5be3dd74ab09bc50",
     "locked": true,
     "schema_version": 3,
     "solution": false,
     "task": false
    }
   },
   "source": [
    "In this exercise we will investigate two state-of-the-art algorithms for handling a continuous control set: deep deterministic policy gradient (DDPG) and proximal policy optimization (PPO).\n",
    "\n",
    "We will examine their performance on [Goddard's rocket problem](https://github.com/osannolik/gym-goddard).\n",
    "This environment comes prepackaged in this notebook's folder, so it can be easily imported.\n",
    "\n",
    "```\n",
    "First formulated by R. H. Goddard around 1910, this is a classical problem within dynamic optimization and optimal control. The task is simply to find the optimal thrust profile for a vertically ascending rocket in order for it to reach the maximum possible altitude, given that its mass decreases as the fuel is spent and that it is subject to varying drag and gravity.\n",
    "```\n",
    "\n",
    "The gym's observation space is the rocket's vertical position, velocity and mass, which are already normalized to the range $[0, 1]$.\n",
    "\n",
    "The rocket engine is assumed to be throttled such that the thrust can be continuously controlled between 0 up to its maximum limit, which translates to an action space $\\mathcal U \\in [0, 1]$.\n",
    "\n",
    "During the episode the reward is set to $r=0$.\n",
    "When the episode is terminating, the (normalized) maximum altitude reached by the rocket during the experiment is used as reward.\n",
    "If the rocket has not depleted its tank after $300$ steps, a reward of $r=-1$ is applied and the episode terminates.\n",
    "\n",
    "![](rocket_spacex.jpg)\n",
    "(Photo by <a href=\"https://unsplash.com/@spacex?utm_source=unsplash&utm_medium=referral&utm_content=creditCopyText\">SpaceX</a> on <a href=\"https://unsplash.com/s/photos/rocket?utm_source=unsplash&utm_medium=referral&utm_content=creditCopyText\">Unsplash</a>)\n",
    "  \n",
    "\n",
    "In order for the full notebook to run through, you will also need the stable-baslines3 package.\n",
    "\n",
    "```\n",
    "pip install stable-baselines3[extra]\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0dde9aac",
   "metadata": {
    "nbgrader": {
     "grade": false,
     "grade_id": "cell-f7ba0aeae957d480",
     "locked": true,
     "schema_version": 3,
     "solution": false,
     "task": false
    }
   },
   "outputs": [],
   "source": [
    "from rocket_env import GoddardEnv\n",
    "from tqdm import tqdm\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "from torch.optim import Adam\n",
    "from stable_baselines3 import PPO, DDPG\n",
    "from stable_baselines3.common.callbacks import BaseCallback\n",
    "from stable_baselines3.common.monitor import Monitor\n",
    "from stable_baselines3.common.noise import NormalActionNoise\n",
    "from copy import deepcopy\n",
    "\n",
    "OPTIMAL_CONTROL = 0.0122079818367078\n",
    "RANDOM_AVG = 0.00995\n",
    "RANDOM_BEST = 0.01148"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6b84bfa8",
   "metadata": {
    "nbgrader": {
     "grade": false,
     "grade_id": "cell-d77d2440cb4fa969",
     "locked": true,
     "schema_version": 3,
     "solution": false,
     "task": false
    }
   },
   "source": [
    "## 1) DDPG"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3ec258ce",
   "metadata": {
    "nbgrader": {
     "grade": false,
     "grade_id": "cell-027671c6914671f2",
     "locked": true,
     "schema_version": 3,
     "solution": false,
     "task": false
    }
   },
   "source": [
    "Deep Deterministic Policy Gradient (DDPG) was first introduced [2015](https://arxiv.org/abs/1509.02971).\n",
    "It uses four neural networks to operate on the continuous state and action space.\n",
    "The first two we already know from the last exercise introduced Actor-Critic method: a critic to estimate the q-value function and a (now deterministic) actor which - once completely trained - should give us the best action for a specific state with respect to the q-value function.\n",
    "Additionally, DDPG provides two target networks which are (in the beginning) copies of the actor and critic. These target networks are updated time-delayed in a low-pass filter manner.\n",
    "This enhances stability during the learning process because the networks used to calculate the target are not directly updated on the fly.\n",
    "For more information see the example [here](https://spinningup.openai.com/en/latest/algorithms/ddpg.html).\n",
    "\n",
    "During the learning process, the critic is updated based on minimizing (!) the following loss function (```critic_loss```):\n",
    "\n",
    "\\begin{equation}\n",
    "\\mathcal L(w, {\\mathcal D}) = \\underset{(x,u,r,x',d) \\sim {\\mathcal D}}{{\\mathrm E}}\\left[\n",
    "    \\Bigg(  \\left(r + \\gamma \\hat{q}_{}(x',\\mu(x', \\theta^-), w^-) \\right) - \\hat{q}(x,u, w) \\Bigg)^2\n",
    "    \\right].\n",
    "\\end{equation}   \n",
    "\n",
    "Here, the target is calculated using the target networks. $w$ and $w^-$ are the parameters of the critic and critic-target networks $\\hat{q}$, respectively $\\theta$ and $\\theta^-$ define the parameters of the actor and actor-target networks.\n",
    "\n",
    "The policy / actor network is updated to maximize (!) the expected return (```actor_loss```)\n",
    "\n",
    "\\begin{align}\n",
    "    \\max_{\\theta} \\underset{x \\sim {\\mathcal D}}{{\\mathrm E}}\\left[\\hat{q}(x, \\mu_{}(x, \\theta), w) \\right],\n",
    "\\end{align}  \n",
    "\n",
    "where $\\theta$ is the set of parameters of the policy network $\\mu$.\n",
    "\n",
    "The updates are performed off-policy by sampling from an experience replay buffer ${\\mathcal D}$.\n",
    "Since we are dealing with a deterministic actor, exploration during training must be explicitly considered by the user. It is usually achieved by adding noise to the sampled actions.\n",
    "\n",
    "### Task: Implement the DDPG with PyTorch\n",
    "Write a DDPG algorithm using the Algo. 12.2 in lecture 12! Figure 12.2 might also be helpful to understand the structure of the algorithm.\n",
    "A lot of the functions and classes are already defined to help you out."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bb6c52d9",
   "metadata": {
    "nbgrader": {
     "grade": false,
     "grade_id": "cell-947a0cd7668e21df",
     "locked": true,
     "schema_version": 3,
     "solution": false,
     "task": false
    }
   },
   "source": [
    "Execute the following cell to make use of the defined multi layer perceptron and the plot function!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c36984d1",
   "metadata": {
    "nbgrader": {
     "grade": false,
     "grade_id": "cell-8303cf3b2bb624b8",
     "locked": true,
     "schema_version": 3,
     "solution": false,
     "task": false
    }
   },
   "outputs": [],
   "source": [
    "def mlp(sizes, activation, output_activation=nn.Identity):\n",
    "    \"\"\"Defines a multi layer perceptron using pytorch layers and activation funtions.\n",
    "\n",
    "    Args:\n",
    "        sizes: The sizes of all the network layers. This includes the input and output layers.\n",
    "        activation: The activation function for the hidden layers\n",
    "        output_action: The activation function for the output layer\n",
    "\n",
    "    Returns:\n",
    "        The multilayer-perceptron PyTorch model\n",
    "    \"\"\"\n",
    "    layers = []\n",
    "    for j in range(len(sizes)-1):\n",
    "        act = activation if j < len(sizes)-2 else output_activation\n",
    "        layers += [nn.Linear(sizes[j], sizes[j+1]), act()]\n",
    "    return nn.Sequential(*layers)\n",
    "\n",
    "def plot_reward_trends(logs):\n",
    "    \"\"\"Plots the reward curves for different training runs based on the trainings logs.\"\"\"\n",
    "    repeats = len(logs['envs'])\n",
    "    ncols = 2\n",
    "    nrows = int(np.ceil(repeats / ncols))\n",
    "    fig, axes = plt.subplots(nrows=nrows, ncols=ncols, figsize=(ncols*8, nrows*4), sharex=True, sharey=True)\n",
    "    for ax, _env, _model in zip(axes.flatten(), logs['envs'], logs['models']):\n",
    "        ax.plot(_env.get_episode_rewards(), label='rewards', color='green')\n",
    "\n",
    "        ax.set_ylabel('Reward')\n",
    "        ax.set_xlabel('Episode')\n",
    "        ax.set_ylim(0, 0.015)\n",
    "        ax.axhline(OPTIMAL_CONTROL, ls='--', color='red', label='optimal')\n",
    "        ax.axhline(RANDOM_BEST, ls='--', color='orange', label='best rnd actions')\n",
    "        ax.axhline(RANDOM_AVG, ls='--', color='yellow', label='average rnd actions')\n",
    "        ax.legend(loc='lower left')\n",
    "        ax2 = ax.twinx()\n",
    "        ax2.plot(_env.get_episode_lengths(), label='ep lengths')\n",
    "        ax2.legend()\n",
    "        ax2.set_ylabel('Steps')\n",
    "    fig.tight_layout()\n",
    "\n",
    "def test_agent(logs):\n",
    "    \"\"\"Deterministic test of the agents based on the trainings logs.\"\"\"    \n",
    "    num_test_episodes = 400\n",
    "    for i, (env, agent) in enumerate(zip(logs['envs'], logs['models'])):\n",
    "        episode_reward = 0\n",
    "        episode_len = 0\n",
    "        state, info = env.reset()\n",
    "        for j in tqdm(range(num_test_episodes)):\n",
    "            episode_len += 1\n",
    "            action = agent.policy(state, deterministic=True)\n",
    "            next_state, reward, terminated, truncated, _ = env.step(action)\n",
    "            episode_reward += reward\n",
    "            state = next_state\n",
    "            if terminated or truncated:\n",
    "                state, info = env.reset()\n",
    "                break\n",
    "\n",
    "        print(f'Reward during test ({i}): {episode_reward}')\n",
    "        env.close()\n",
    "    print('Optimal control: 0.0122')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "277fc242",
   "metadata": {
    "nbgrader": {
     "grade": false,
     "grade_id": "cell-b4bd4f5ae4c71df8",
     "locked": true,
     "schema_version": 3,
     "solution": false,
     "task": false
    }
   },
   "source": [
    "Execute the following cells to make use of the predefined actor & critic and the replay buffer!\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f1b2a7a8",
   "metadata": {
    "nbgrader": {
     "grade": false,
     "grade_id": "cell-93290e96cc879232",
     "locked": true,
     "schema_version": 3,
     "solution": false,
     "task": false
    }
   },
   "outputs": [],
   "source": [
    "class Actor(nn.Module):\n",
    "    \"\"\"The actor network is used to decide on actions based on the current state. This class\n",
    "    implements two functions for this 'forward' and 'act'. One of these is used to interact\n",
    "    with the environment and one is used to evaluate to criticise the actor. Consider which is\n",
    "    necessary in which of these cases.\"\"\"\n",
    "\n",
    "    def __init__(self, obs_dim, act_dim, hidden_sizes, activation=nn.ReLU, act_limit=1):\n",
    "        super().__init__()\n",
    "        mu_sizes = [obs_dim] + list(hidden_sizes) + [act_dim]\n",
    "        self.mu = mlp(mu_sizes, activation, nn.Sigmoid)\n",
    "        self.act_limit = act_limit\n",
    "\n",
    "    def forward(self, state):\n",
    "        \"\"\"Decides on an action given a state. Note that by calling forward the computation is \n",
    "        taken into account for backpropagation.\"\"\"\n",
    "        return self.act_limit * self.mu(state)\n",
    "\n",
    "    def act(self, state):\n",
    "        \"\"\"Decides on an action given a state. Note that this computation is not taken into\n",
    "        account for backpopagation because of the 'torch.no_grad'.\"\"\"\n",
    "        with torch.no_grad():\n",
    "            return self.act_limit * self.mu(state).numpy()\n",
    "    \n",
    "\n",
    "class Critic(nn.Module):\n",
    "    \"\"\"The critic network is used to estimate the q-value for a given state-action pair.\"\"\"\n",
    "\n",
    "    def __init__(self, obs_dim, act_dim, hidden_sizes, activation=nn.ReLU):\n",
    "        super().__init__()\n",
    "        self.q = mlp([obs_dim + act_dim] + list(hidden_sizes) + [1], activation)\n",
    "\n",
    "    def forward(self, state, action):\n",
    "        \"\"\"Estimates the q-value for the given state-action pair.\"\"\"\n",
    "\n",
    "        # torch.cat concatenates action to state\n",
    "        q = self.q(torch.cat([state, action], dim=-1))\n",
    "        return torch.squeeze(q, -1)  # To ensure q has right shape.\n",
    "    \n",
    "class ReplayBuffer:\n",
    "    \"\"\"The replay buffer acts as the memory of the algorithm. When the buffer is fully filled,\n",
    "    the oldest data is replaced with new incoming data.\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, obs_dim, action_dim, buffer_size):\n",
    "        \n",
    "        self.state_buf = np.zeros((buffer_size, obs_dim), dtype=np.float32)\n",
    "        self.next_state_buf = np.zeros((buffer_size, obs_dim), dtype=np.float32)\n",
    "        self.action_buf = np.zeros((buffer_size, action_dim), dtype=np.float32)\n",
    "        self.reward_buf = np.zeros(buffer_size, dtype=np.float32)\n",
    "        self.terminated_buf = np.zeros(buffer_size, dtype=np.float32)\n",
    "        self.truncated_buf = np.zeros(buffer_size, dtype=np.float32)\n",
    "        self.ptr, self.size, self.max_size = 0, 0, buffer_size\n",
    "\n",
    "    def push(self, state, action, reward, next_state, terminated, truncated):\n",
    "        \"\"\"Adds a new data point to the memory. For the case where the buffer is already full.\n",
    "        The oldest data is replaced with the new data.\n",
    "        \"\"\"\n",
    "        self.state_buf[self.ptr] = state\n",
    "        self.next_state_buf[self.ptr] = next_state\n",
    "        self.action_buf[self.ptr] = action\n",
    "        self.reward_buf[self.ptr] = reward\n",
    "        self.terminated_buf[self.ptr] = terminated\n",
    "        self.truncated_buf[self.ptr] = truncated\n",
    "        self.ptr = (self.ptr + 1) % self.max_size\n",
    "        self.size = min(self.size + 1, self.max_size)\n",
    "\n",
    "    def fetch(self, batch_size=32):\n",
    "        \"\"\"Fetches a random batch from the memory.\"\"\"\n",
    "        \n",
    "        idxs = np.random.randint(0, self.size, size=batch_size)\n",
    "        \n",
    "        batch = dict(state=self.state_buf[idxs],\n",
    "                     next_state=self.next_state_buf[idxs],\n",
    "                     action=self.action_buf[idxs],\n",
    "                     reward=self.reward_buf[idxs],\n",
    "                     terminated=self.terminated_buf[idxs],\n",
    "                     truncated=self.truncated_buf[idxs]\n",
    "                )\n",
    "        \n",
    "        return {k: torch.as_tensor(v, dtype=torch.float32) for k, v in batch.items()}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "19174417",
   "metadata": {
    "nbgrader": {
     "grade": false,
     "grade_id": "cell-2639a410b1080568",
     "locked": true,
     "schema_version": 3,
     "solution": false,
     "task": false
    }
   },
   "source": [
    "Here, fill in the following code template to write a DDPG agent!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9c9a893e",
   "metadata": {
    "nbgrader": {
     "grade": false,
     "grade_id": "cell-fbb8e1d539cd0caf",
     "locked": false,
     "schema_version": 3,
     "solution": true,
     "task": false
    }
   },
   "outputs": [],
   "source": [
    "class DDPG_agent:\n",
    "    \"\"\"Reference: https://github.com/openai/spinningup/blob/master/spinup/algos/pytorch/ddpg/ddpg.py\n",
    "\n",
    "    This class takes care of handling the networks and almost all of the heavy lifting for the algorithm.\n",
    "    \"\"\"    \n",
    "    \n",
    "    def __init__(self, env, actor_hidden_size, actor_number_layers, critic_hidden_size, critic_number_layers,\n",
    "                 buffer_size, actor_lr, critic_lr, gamma, batch_size, learning_starts, tau):\n",
    "        self.actor_lr = actor_lr\n",
    "        self.critic_lr = critic_lr\n",
    "        self.batch_size = batch_size\n",
    "        self.learning_starts = learning_starts\n",
    "        self.gamma = gamma\n",
    "        self.tau = tau\n",
    "\n",
    "        obs_dim = env.observation_space.shape[0]\n",
    "        act_dim = env.action_space.shape[0]\n",
    "    \n",
    "        # initialize actor, critic and the replaybuffer into the variables given below\n",
    "        # consider which layer sizes are necessary for which network\n",
    "        self.actor = None\n",
    "        self.critic = None\n",
    "        self.replay_buffer = None\n",
    "        \n",
    "        # YOUR CODE HERE\n",
    "        raise NotImplementedError()\n",
    "        \n",
    "        # defines target networks\n",
    "        self.actor_target = deepcopy(self.actor)\n",
    "        self.critic_target = deepcopy(self.critic)\n",
    "        \n",
    "        # Uses Adam optimizer (see to ex11 for more explanation)\n",
    "        self.actor_optimizer = Adam(self.actor.parameters(), lr=actor_lr)\n",
    "        self.critic_optimizer = Adam(self.critic.parameters(), lr=critic_lr)\n",
    "\n",
    "        # define action noise space for clipping\n",
    "        self.action_space_high = env.action_space.high\n",
    "        self.action_space_low = env.action_space.low\n",
    "\n",
    "    def critic_loss(self, data):\n",
    "        \"\"\"Compute the loss for the critic network as given above or in Algo. 12.2. Remember resetting\n",
    "        the computed gradients from the last iteration to zero. Compute the target with bootstrapping\n",
    "        and compare it to the current estimate of the critic.\n",
    "\n",
    "        **Hint**: You want to compute the gradient of the loss only w.r.t. the parameters of one specific\n",
    "        network. Which network is that and why? \n",
    "        Use ```with torch.no_grad():``` for any operation with any other network.\n",
    "        \n",
    "        Args:\n",
    "            data: A batch of data sampled from the memory. The unwrapping of the data batch is given below.\n",
    "\n",
    "        Returns:\n",
    "            q_loss: The critic loss with attached computational graph\n",
    "        \"\"\"\n",
    "        state, action, rerwad, next_state, terminated, truncated = data['state'], data['action'], data['reward'], data['next_state'], data['terminated'], data['truncated']\n",
    "        \n",
    "        # YOUR CODE HERE\n",
    "        raise NotImplementedError()\n",
    "        return q_loss\n",
    "    \n",
    "    def actor_loss(self, data):\n",
    "        \"\"\"Compute the loss for the actor network as given above or in Algo. 12.2. Remember resetting\n",
    "        the computed gradients from the last iteration to zero.\n",
    "\n",
    "        **Hint**: Note that this is a gradient ascent(!) step. What needs to be changed as a result?\n",
    "\n",
    "        Args:\n",
    "            data: A batch of data sampled from the memory. Only the state is actually needed\n",
    "\n",
    "        Returns:\n",
    "            loss: The actor loss with attached computational graph\n",
    "        \"\"\"\n",
    "        state = data['state']\n",
    "        \n",
    "        # YOUR CODE HERE\n",
    "        raise NotImplementedError()\n",
    "        return loss\n",
    "        \n",
    "    \n",
    "    def learn(self, number_updates_per_step=1):\n",
    "        \"\"\"\n",
    "        Fetches number_updates_per_step-times from replay_buffer, computes targets and calculates losses to \n",
    "        update the q-function using gradient descent and the policy function using gradient ascent\n",
    "        In the end updates the target networks\n",
    "        \"\"\"\n",
    "\n",
    "        for _ in range(number_updates_per_step):\n",
    "\n",
    "            data = self.replay_buffer.fetch(self.batch_size)\n",
    "\n",
    "            # Gradient descent for Q-function/critic\n",
    "            # YOUR CODE HERE\n",
    "            raise NotImplementedError()\n",
    "\n",
    "            # Freeze Q-network so you don't waste computational effort\n",
    "            # computing gradients for it during the policy learning step.\n",
    "            for p in self.critic.parameters():\n",
    "                p.requires_grad = False\n",
    "\n",
    "            # Gradient ascent step for mu/actor\n",
    "            # YOUR CODE HERE\n",
    "            raise NotImplementedError()\n",
    "\n",
    "            # Unfreeze Q-network so you can optimize it at next DDPG step.\n",
    "            for p in self.critic.parameters():\n",
    "                p.requires_grad = True\n",
    "\n",
    "            # Finally, update target networks by polyak averaging.\n",
    "            with torch.no_grad():\n",
    "                # update Critc-target\n",
    "                for p, p_targ in zip(self.critic.parameters(), self.critic_target.parameters()):\n",
    "                    # uses in-place operations \"mul_\", \"add_\" to update target\n",
    "                    # params, as opposed to \"mul\" and \"add\", which would make new tensors.\n",
    "                    p_targ.data.mul_(1 - self.tau)\n",
    "                    p_targ.data.add_((self.tau) * p.data)\n",
    "\n",
    "                # update Actor-target\n",
    "                for p, p_targ in zip(self.actor.parameters(), self.actor_target.parameters()):\n",
    "                    # uses in-place operations \"mul_\", \"add_\" to update target\n",
    "                    # params, as opposed to \"mul\" and \"add\", which would make new tensors.\n",
    "                    p_targ.data.mul_(1 - self.tau)\n",
    "                    p_targ.data.add_((self.tau) * p.data)\n",
    "            \n",
    "    def policy(self, state, deterministic=False, noise_scale=0.1):\n",
    "        \"\"\"Decides on an action, given the current state. Adds scaled noise if the choice should\n",
    "        not be deterministic and clips the action depending on the action space.\n",
    "\n",
    "        Args:\n",
    "            state: The current state\n",
    "            deterministic: Whether to choose the action deterministically or add noise to the choice\n",
    "\n",
    "        Returns:\n",
    "            action: The chosen action\n",
    "        \"\"\"\n",
    "        # YOUR CODE HERE\n",
    "        raise NotImplementedError()\n",
    "        return action      "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5f947002",
   "metadata": {
    "nbgrader": {
     "grade": false,
     "grade_id": "cell-f1983e65eedd72ba",
     "locked": true,
     "schema_version": 3,
     "solution": false,
     "task": false
    }
   },
   "source": [
    "Run the following cell to train your agent. Note that, in contrast to classic RL applications, in standard libraries you have often the possibility to update the agent after full episodes instead of stepwise.\n",
    "This could be helpful in our application due to the sparse reward signal of the rocket environment instead of after each environment step.\n",
    "To overcome this issue we introduce the variable `learning_starts` and run 500 steps in the environment at first, to be sure that we have at least one full episode in the replay buffer to learn from.\n",
    "\n",
    "To compare episode-wise vs step-wise updating, both is executed 4 times in the following cell.\n",
    "The learning could take some minutes. If you want to reduce it you can do only 1 run for each update for example (reduce `repeats`) or run less steps at all (reduce `total_timesteps`)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9bb2e05b",
   "metadata": {
    "nbgrader": {
     "grade": false,
     "grade_id": "cell-c7977ed401a0f23b",
     "locked": true,
     "schema_version": 3,
     "solution": false,
     "task": false
    }
   },
   "outputs": [],
   "source": [
    "train_logs_ddpg = {'episode': {'envs': [], 'models': []}, 'step_wise': {'envs': [], 'models': []}}\n",
    "repeats = 4\n",
    "\n",
    "total_timesteps = 20000\n",
    "\n",
    "\n",
    "for train_freq, logs_ddpg in train_logs_ddpg.items():\n",
    "\n",
    "    for rep in range(repeats):\n",
    "\n",
    "\n",
    "        env = Monitor(GoddardEnv(render_mode=None))\n",
    "\n",
    "        myDDPG_agent = DDPG_agent(env=env, actor_hidden_size=8, actor_number_layers=1,\n",
    "                               critic_hidden_size=8, critic_number_layers=1, buffer_size=int(1e6),\n",
    "                               actor_lr=1e-4, critic_lr=1e-4, gamma=0.999, batch_size=256, learning_starts=500, tau=0.005)\n",
    "\n",
    "        number_updates_per_step = 0\n",
    "        episode_reward = 0\n",
    "        episode_len = 0\n",
    "        rewards = []\n",
    "        episode_len_vec = []\n",
    "\n",
    "        state, info = env.reset()\n",
    "\n",
    "        for j in tqdm(range(total_timesteps)):\n",
    "            episode_len += 1\n",
    "            number_updates_per_step += 1\n",
    "\n",
    "            action = myDDPG_agent.policy(state)\n",
    "            next_state, reward, terminated, truncated, _ = env.step(action)\n",
    "            episode_reward += reward\n",
    "            myDDPG_agent.replay_buffer.push(state, action, reward, next_state, terminated, truncated)\n",
    "\n",
    "            if train_freq == 'step_wise' and j > myDDPG_agent.learning_starts:\n",
    "                # Train every step using myDDPG_agent.batch_size batch\n",
    "                myDDPG_agent.learn()\n",
    "            \n",
    "            state = next_state\n",
    "\n",
    "            if truncated or terminated:\n",
    "                if train_freq == 'episode' and j > myDDPG_agent.learning_starts:\n",
    "                # Train after each episode number_updates_per_step times using myDDPG_agent.batch_size batch\n",
    "                    myDDPG_agent.learn(number_updates_per_step)\n",
    "                state, info = env.reset()\n",
    "                rewards.append(episode_reward)\n",
    "                episode_len_vec.append(episode_len)\n",
    "                episode_reward = 0\n",
    "                episode_len = 0\n",
    "                number_updates_per_step = 0\n",
    "\n",
    "        logs_ddpg['envs'] += [env]\n",
    "        logs_ddpg['models'] += [myDDPG_agent]\n",
    "        env.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4280d889",
   "metadata": {
    "nbgrader": {
     "grade": false,
     "grade_id": "cell-f63074dfa3920133",
     "locked": true,
     "schema_version": 3,
     "solution": false,
     "task": false
    }
   },
   "outputs": [],
   "source": [
    "plot_reward_trends(train_logs_ddpg['episode'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "33d9dfc7",
   "metadata": {
    "nbgrader": {
     "grade": false,
     "grade_id": "cell-ddffaf6c11d49a3a",
     "locked": true,
     "schema_version": 3,
     "solution": false,
     "task": false
    }
   },
   "outputs": [],
   "source": [
    "plot_reward_trends(train_logs_ddpg['step_wise'])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7d583148",
   "metadata": {
    "nbgrader": {
     "grade": false,
     "grade_id": "cell-04b7fa3a7ef613f5",
     "locked": true,
     "schema_version": 3,
     "solution": false,
     "task": false
    }
   },
   "source": [
    "Use the following cell to test your agent on the env using deterministic actions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0bc1f4cc",
   "metadata": {
    "nbgrader": {
     "grade": false,
     "grade_id": "cell-b85d8e3dd9ad383b",
     "locked": true,
     "schema_version": 3,
     "solution": false,
     "task": false
    }
   },
   "outputs": [],
   "source": [
    "test_agent(train_logs_ddpg['episode'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "45af0043",
   "metadata": {
    "nbgrader": {
     "grade": false,
     "grade_id": "cell-94f13e77c62bc467",
     "locked": true,
     "schema_version": 3,
     "solution": false,
     "task": false
    }
   },
   "outputs": [],
   "source": [
    "test_agent(train_logs_ddpg['step_wise'])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1240d5c7",
   "metadata": {
    "nbgrader": {
     "grade": false,
     "grade_id": "cell-57ec48fad41194c8",
     "locked": true,
     "schema_version": 3,
     "solution": false,
     "task": false
    }
   },
   "source": [
    "### Demo: StableBaselines3 usage\n",
    "\n",
    "Alternatively, you can use readily available Python packages such as [stable-baselines3](https://github.com/DLR-RM/stable-baselines3) for employment of state-of-the-art algorithms.\n",
    "In what follows below, the DDPG algorithm as utilized by stable-baselines3 is showcased.\n",
    "\n",
    "Note that in stable baselines the deterministic actor/policy is named $\\pi$ and not $\\mu$ (as in the lecture)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "576cea94",
   "metadata": {
    "nbgrader": {
     "grade": false,
     "grade_id": "cell-4bc64d70ae9709c7",
     "locked": true,
     "schema_version": 3,
     "solution": false,
     "task": false
    }
   },
   "outputs": [],
   "source": [
    "total_timesteps = int(20e3)\n",
    "\n",
    "\n",
    "class Pbar(BaseCallback):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        self.pbar = tqdm(desc='Training', total=total_timesteps, position=0, leave=True)\n",
    "\n",
    "    def _on_rollout_end(self) -> None:\n",
    "        \"\"\"This event is triggered before updating the policy.\"\"\"\n",
    "        pass\n",
    "\n",
    "    def _on_step(self) -> bool:\n",
    "        self.pbar.update()\n",
    "        return True\n",
    "    \n",
    "    def _on_training_end(self) -> None:\n",
    "        \"\"\"This event is triggered before exiting the `learn()` method.\"\"\"\n",
    "        print(\"\")\n",
    "    \n",
    "policy_kwargs = dict(activation_fn=torch.nn.ReLU, net_arch=[dict(pi=[8], vf=[8])])\n",
    "\n",
    "repeats = 4\n",
    "logs = {'envs': [], 'models': []}\n",
    "for _ in range(repeats):\n",
    "    \n",
    "    env = Monitor(GoddardEnv(render_mode=None))\n",
    "    \n",
    "    n_actions = env.action_space.shape[-1]\n",
    "    action_noise = NormalActionNoise(mean=np.zeros(n_actions), sigma=0.1 * np.ones(n_actions))\n",
    "    \n",
    "    policy_kwargs = dict(activation_fn=torch.nn.ReLU, net_arch=dict(pi=[8] * 1, qf=[8] * 1))\n",
    "    model = DDPG(\"MlpPolicy\", env, policy_kwargs=policy_kwargs, batch_size=256, device='cpu',\n",
    "             learning_rate=1e-4, gamma=0.999, learning_starts=500)\n",
    "    model.learn(total_timesteps=total_timesteps, callback=Pbar())\n",
    "    logs['envs'] += [env]\n",
    "    logs['models'] += [model]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e1bf745c",
   "metadata": {
    "nbgrader": {
     "grade": false,
     "grade_id": "cell-354440d67e2e0b95",
     "locked": true,
     "schema_version": 3,
     "solution": false,
     "task": false
    }
   },
   "outputs": [],
   "source": [
    "plot_reward_trends(logs)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0e84675a",
   "metadata": {
    "nbgrader": {
     "grade": false,
     "grade_id": "cell-f7ad3146465beb22",
     "locked": true,
     "schema_version": 3,
     "solution": false,
     "task": false
    }
   },
   "source": [
    "Use the following cell to execute the last trained agent"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bc0cddd2",
   "metadata": {
    "nbgrader": {
     "grade": false,
     "grade_id": "cell-c25a8b2418b0a9a8",
     "locked": true,
     "schema_version": 3,
     "solution": false,
     "task": false
    }
   },
   "outputs": [],
   "source": [
    "n_episodes = 1\n",
    "max_steps_per_episode = 500\n",
    "tst_logs = {'rewards': []}\n",
    "\n",
    "env.env.render_mode = \"human\"\n",
    "\n",
    "\n",
    "for ep in range(n_episodes):\n",
    "    state, info = env.reset()\n",
    "    done = np.zeros(1, dtype=np.bool_)\n",
    "    cum_rew = 0\n",
    "    k = 0\n",
    "    while not np.all(done) and k < max_steps_per_episode:\n",
    "        action, _ = model.predict(state, deterministic=True) \n",
    "        state, reward, truncated, terminated, info = env.step(action)\n",
    "        done = truncated or terminated\n",
    "        cum_rew += reward\n",
    "        k += 1\n",
    "\n",
    "    tst_logs['rewards'].append(cum_rew)\n",
    "rewards = np.array(tst_logs['rewards']).ravel()\n",
    "env.close()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8e03bcfb",
   "metadata": {
    "nbgrader": {
     "grade": false,
     "grade_id": "cell-038ec2a3065ebfd0",
     "locked": true,
     "schema_version": 3,
     "solution": false,
     "task": false
    }
   },
   "source": [
    "## 2) PPO\n",
    "\n",
    "The [original paper from 2017](https://arxiv.org/abs/1707.06347) for the PPO came up with an idea to combine A3C (having multiple workers) and TRPO (using a trust region to improve the actor).\n",
    "The PPO algorithm achieves this by clipping gradients in order to ensure that the new policy is not too far away from the old one.\n",
    "\n",
    "In contrast to DDPG, PPO is an on-policy algorithm. In order to still apply mini-batch training, there is a so-called roll-out-buffer that is filled up with the current policy, on whose base a gradient ascent update would be done.\n",
    "\n",
    "There are two variants on the PPO algorithm, from which we will implement the clip variant.\n",
    "\n",
    "The actor (policy) update is computed according to (```actor_loss```)\n",
    "\n",
    "\\begin{align}\n",
    "\\theta_{k+1} &= \\arg \\max_{\\theta} \\underset{\\langle x,u \\rangle \\sim \\pi_{\\theta_k}}{{\\mathbb E}}\\left[\\mathcal L(X,U,\\theta_k, \\theta)\\right], \\\\\n",
    "\\mathcal L(x,u,\\theta_k,\\theta) &= \\min\\left(\n",
    "\\frac{\\pi_{\\theta}(u|x)}{\\pi_{\\theta_k}(u|x)}  A^{\\pi_{\\theta_k}}(x,u), \\;\\;\n",
    "\\text{clip}\\left(\\frac{\\pi_{\\theta}(u|x)}{\\pi_{\\theta_k}(u|x)}, 1 - \\epsilon, 1+\\epsilon \\right) A^{\\pi_{\\theta_k}}(x, u)\n",
    "\\right),\n",
    "\\end{align}\n",
    "where the expectation operator denotes the empirical average across the roll-out-buffer, $\\theta$ the possible next weights of the actor, $\\epsilon$ denotes the threshold within the next update is allowed, and $\\mathcal L$ the long-term return.\n",
    "\n",
    "On the opposite side, the critic's weights $\\omega_k$ are updated through the plain stochastic mini-batch gradient descent on the mean squared error $\\mathcal C$ between immediately seen rewards $r_{k+1}$ and corresponding estimated value $v_{w_k}(x_k)$ across the roll-out-buffer (```critic_loss```):\n",
    "\\begin{align}\n",
    "w_{k+1} &\\leftarrow \\omega_k - \\alpha \\nabla\\mathcal C(X,R_{k+1},\\omega_k), \\\\\n",
    "\\mathcal C(x,r_{k+1}, \\omega_k) &= \\left(v_{\\omega_k} (x_k)- r_{k+1}\\right)^2\n",
    "\\end{align}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e1c51333",
   "metadata": {
    "nbgrader": {
     "grade": false,
     "grade_id": "cell-37e0d6f3e79fb7f6",
     "locked": true,
     "schema_version": 3,
     "solution": false,
     "task": false
    }
   },
   "source": [
    "### Task: Implement the PPO clip variant with PyTorch\n",
    "Fill in the below code to run the PPO-clip variant on the rocket environment. Use the Algo. 13.1 from lecture 13 for guidance.\n",
    "\n",
    "For simplicity, do not contemplate a vectorized environment for synchronous training (like in A2C or A3C).\n",
    "Moreover, calculate the generalized advantage estimate with $\\lambda = 1$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "641b1e59",
   "metadata": {
    "nbgrader": {
     "grade": false,
     "grade_id": "cell-bfaedbfa4a92491d",
     "locked": false,
     "schema_version": 3,
     "solution": true,
     "task": false
    }
   },
   "outputs": [],
   "source": [
    "class StochActor(nn.Module):\n",
    "    \"\"\"This stochastic actor learns a Gaussian distribution. While the mean value \n",
    "    is learnt by a full-fledged MLP, the standard deviation is denoted by a\n",
    "    single trainable weight. With Pytorch's distribution package, probabilities \n",
    "    given a certain distribution and an action can be calculated.\"\"\"\n",
    "    def __init__(self, state_dim, action_dim):\n",
    "        super().__init__()\n",
    "        # The standard deviation is just one trainable weight\n",
    "        self.log_std = torch.nn.Parameter(\n",
    "                        torch.as_tensor(\n",
    "                            -.5 * np.ones(action_dim, dtype=np.float32)))\n",
    "        # the mean value is estimated by a full MLP\n",
    "        self.mu_net = mlp(list(state_dim) + [8] + list(action_dim), nn.Tanh, nn.Sigmoid)\n",
    "        \n",
    "    def _distribution(self, state):\n",
    "        \"\"\"Returns a gaussian distribution where mean and standard deviation are chosen in\n",
    "        dependence of the current state.\n",
    "        \"\"\"\n",
    "        return torch.distributions.normal.Normal(\n",
    "                self.mu_net(state), torch.exp(self.log_std))\n",
    "    \n",
    "    def forward(self, state, action=None):\n",
    "        \"\"\"Computes the action distribtuion for the current state and the probability of an action\n",
    "        in that policy.\n",
    "\n",
    "        Args:\n",
    "            state: The state for which to compute the action distribtuion\n",
    "            action: An action for which to compute the log probability\n",
    "\n",
    "        Returns:\n",
    "            pi: The computed action distribution\n",
    "            logp_a: The log probability for the given action in the computed distribution\n",
    "        \"\"\"\n",
    "        pi = self._distribution(state)\n",
    "        # if action is None, logp_a will be, too\n",
    "        if action is None:\n",
    "            logp_a = None\n",
    "        else:\n",
    "            logp_a = pi.log_prob(action).sum(axis=-1)\n",
    "        return pi, logp_a\n",
    "\n",
    "\n",
    "class RolloutBuffer:\n",
    "    \"\"\"The rollout buffer serves as the memory of this algorithm. Note while it looks similar to the\n",
    "    replay buffer, it fullfils a very different role:\n",
    "\n",
    "    The replay buffer (in DDPG) is filled more or less independent of the training. It just gathers \n",
    "    the data that is put into it and releases old data when there is no more space. Fetching a\n",
    "    batch from this buffer just gives you a random batch of data tuples from interactions that were\n",
    "    done at some point during the training.\n",
    "    \n",
    "    The rollout buffer (in PPO) is filled with a full episode that is then used for learning (this\n",
    "    is more akin to Monte Carlo simulations from e.g. exercise 4). This means that the data batch\n",
    "    from the rollout buffer is always the sequence of interactions from the last episode.\n",
    "    \"\"\"\n",
    "    def __init__(self, size, action_dim, state_dim):\n",
    "        self.action_buf = np.zeros((size, action_dim[0]), dtype=np.float32)\n",
    "        self.state_buf = np.zeros((size, state_dim[0]), dtype=np.float32)\n",
    "        self.rew_buf = np.zeros(size+1, dtype=np.float32)\n",
    "        self.val_buf = np.zeros(size+1, dtype=np.float32)\n",
    "        self.logp_buf = np.zeros(size, dtype=np.float32)\n",
    "        self.i = 0\n",
    "        self.max_size = size\n",
    "        \n",
    "    def push(self, state, action, reward, value, logp):\n",
    "        \"\"\"Append a sample to the buffer.\"\"\"\n",
    "        self.state_buf[self.i] = state\n",
    "        self.action_buf[self.i] = action\n",
    "        self.rew_buf[self.i] = reward\n",
    "        self.val_buf[self.i] = value\n",
    "        self.logp_buf[self.i] = logp\n",
    "        self.i += 1\n",
    "        \n",
    "    def fetch(self, last_val=None, gamma=0.999):\n",
    "        \"\"\"Get all data from the rollout buffer.\n",
    "\n",
    "        Args:\n",
    "            last_val: An optional last value to add to reward and value buffers\n",
    "            gamma: The discount factor\n",
    "        \n",
    "        Returns:\n",
    "            A dictionary for state, action, rewards-to-go, advantages, and logp values.\n",
    "            Entries might be of different length across multiple calls to this function as episodes\n",
    "            potentially have different lengths. Rewards-to-go and advantages need to be computed from \n",
    "            rewards, values, and gamma, as outlined in https://arxiv.org/abs/1506.02438 (or as done\n",
    "            in the spinning up implementation). Normalize the advantage batch by standard scaling.\n",
    "        \"\"\"\n",
    "        if last_val is not None:\n",
    "            self.rew_buf[self.i] = last_val\n",
    "            self.val_buf[self.i] = last_val\n",
    "        rewards = self.rew_buf[:self.i+1]\n",
    "        values = self.val_buf[:self.i+1]\n",
    "\n",
    "        \n",
    "        # Advantage calculaton\n",
    "        advantages = None\n",
    "        # YOUR CODE HERE\n",
    "        raise NotImplementedError()\n",
    "        \n",
    "        # reward\n",
    "        rewards_to_go = None\n",
    "        # YOUR CODE HERE\n",
    "        raise NotImplementedError()\n",
    "        \n",
    "        # normalize advantages\n",
    "        # YOUR CODE HERE\n",
    "        raise NotImplementedError()\n",
    "\n",
    "\n",
    "        data =  {'state': self.state_buf[:self.i],\n",
    "                 'action': self.action_buf[:self.i],\n",
    "                 'rewards_to_go': rewards_to_go.copy(),  # make stride positive\n",
    "                 'advantages': advantages, \n",
    "                 'logp': self.logp_buf[:self.i]}\n",
    "        self.i = 0\n",
    "        \n",
    "        return {k: torch.as_tensor(v, dtype=torch.float32) for k, v in data.items()}\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a2d4f71d-b8c3-4a9d-96ec-0fef853d6dd9",
   "metadata": {},
   "outputs": [],
   "source": [
    "class PPOAgent:\n",
    "    \"\"\"Reference:\n",
    "    https://github.com/openai/spinningup/blob/master/spinup/algos/pytorch/ppo/ppo.py\n",
    "\n",
    "    This class takes care of handling the networks and almost all of the heavy lifting for the algorithm.\n",
    "    \"\"\"\n",
    "    def __init__(self, buffer_size, env):\n",
    "        self.buf = RolloutBuffer(\n",
    "            buffer_size,\n",
    "            action_dim=env.action_space.shape,\n",
    "            state_dim=env.observation_space.shape\n",
    "        )\n",
    "        # use the below attributes for the actor-critic agent\n",
    "        self.actor = None\n",
    "        self.critic = None\n",
    "        self.n_epochs = 8\n",
    "        self.clip_ratio = 0.1\n",
    "\n",
    "        # initialize actor, critic and their optimizers\n",
    "        # YOUR CODE HERE\n",
    "        raise NotImplementedError()\n",
    "\n",
    "    def critic_loss(self, data):\n",
    "        \"\"\"Compute the loss for the critic network by comparing the state value estimates\n",
    "        of the critic with the rewards to go. Use the mean squared error.\n",
    "        \"\"\"\n",
    "        state, rewards_to_go = data['state'], data['rewards_to_go']\n",
    "        # YOUR CODE HERE\n",
    "        raise NotImplementedError()\n",
    "        return v_loss\n",
    "    \n",
    "    def actor_loss(self, data):\n",
    "        \"\"\"Compute the loss for the actor network as given on slide 15 in lecture 13.\"\"\"\n",
    "\n",
    "        state, action, advantage, logp_old = [\n",
    "            data[k] for k in ('state', 'action', 'advantages', 'logp')\n",
    "        ]\n",
    "        \n",
    "        # YOUR CODE HERE\n",
    "        raise NotImplementedError()\n",
    "\n",
    "        return loss_pi\n",
    "    \n",
    "    def learn(self, last_value=None):\n",
    "        \"\"\"Fetch the rollout buffer for all samples.\n",
    "        Train the actor and critic for n_epochs by computing their\n",
    "        losses iteratively.\n",
    "        The rewards and values buffer needs a zero appended for terminal states,\n",
    "        or the critic's estimate appended in case of intermediate states \n",
    "        that result from timeout. This is denoted by last_value.\"\"\"\n",
    "        # YOUR CODE HERE\n",
    "        raise NotImplementedError()\n",
    "    \n",
    "    def policy(self, state, deterministic=False):\n",
    "        \"\"\"Decides on an action, given the current state. Either the mean of the\n",
    "        of the policy distribution if chosen deterministically or a sample from\n",
    "        that distribution.\n",
    "\n",
    "        Args:\n",
    "            state: The current state\n",
    "            deterministic: Whether to choose the mean or sample from the\n",
    "                distribution\n",
    "        \n",
    "        Returns:\n",
    "            a: The chosen action as a numpy array\n",
    "            v: The state value for the current state\n",
    "            logp_a: The log probability of the chosen action in the policy\n",
    "        \"\"\"\n",
    "        state = torch.as_tensor(state, dtype=torch.float32)\n",
    "        with torch.no_grad():\n",
    "            # YOUR CODE HERE\n",
    "            raise NotImplementedError()\n",
    "        return a.numpy(), v.numpy(), logp_a.numpy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0c9b71ab",
   "metadata": {
    "nbgrader": {
     "grade": false,
     "grade_id": "cell-4cfa3ab6af86faef",
     "locked": true,
     "schema_version": 3,
     "solution": false,
     "task": false
    }
   },
   "outputs": [],
   "source": [
    "# main training loop\n",
    "tqdm(desc='Training', total=total_timesteps, position=None, leave=True)\n",
    "\n",
    "n_episodes = 250\n",
    "max_steps_per_episode = 500\n",
    "train_logs_ppo_custom = {'envs': [], 'models': []}\n",
    "\n",
    "def train_ppo(repeats):\n",
    "    for rep in range(repeats):\n",
    "        env = Monitor(GoddardEnv(render_mode = None))\n",
    "        agent = PPOAgent(max_steps_per_episode, env)\n",
    "\n",
    "        for ep in tqdm(range(n_episodes), desc='Training'):\n",
    "            state, info = env.reset()\n",
    "            done = False\n",
    "            k = 0\n",
    "            while not done and k < max_steps_per_episode:\n",
    "                action, value, logp_a = agent.policy(state) \n",
    "                next_state, reward, terminated, truncated, info = env.step(action)\n",
    "                done = terminated or truncated\n",
    "\n",
    "                # track experience\n",
    "                agent.buf.push(state, action, reward, value, logp_a)\n",
    "\n",
    "                state = next_state\n",
    "                k += 1\n",
    "\n",
    "            # if episode finished with done signal, the long term value should be 0\n",
    "            #  for this terminal state\n",
    "            if done:\n",
    "                last_value = None\n",
    "            else:\n",
    "                _, last_value, _ = agent.policy(state)\n",
    "            agent.learn(last_value)\n",
    "\n",
    "        train_logs_ppo_custom['envs'] += [env]\n",
    "        train_logs_ppo_custom['models'] += [agent]\n",
    "        env.close()\n",
    "        \n",
    "train_ppo(4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "daf36c34",
   "metadata": {
    "nbgrader": {
     "grade": false,
     "grade_id": "cell-ec0e800aac84b980",
     "locked": true,
     "schema_version": 3,
     "solution": false,
     "task": false
    }
   },
   "outputs": [],
   "source": [
    "plot_reward_trends(train_logs_ppo_custom)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2c5224ae",
   "metadata": {
    "nbgrader": {
     "grade": false,
     "grade_id": "cell-f3ba7a0d376cc7ee",
     "locked": true,
     "schema_version": 3,
     "solution": false,
     "task": false
    }
   },
   "outputs": [],
   "source": [
    "test_agent(train_logs_ppo_custom)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "707f9f93",
   "metadata": {
    "nbgrader": {
     "grade": false,
     "grade_id": "cell-94626f12d2902d40",
     "locked": true,
     "schema_version": 3,
     "solution": false,
     "task": false
    }
   },
   "outputs": [],
   "source": [
    "# more samples\n",
    "train_ppo(4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2c86fac7",
   "metadata": {
    "nbgrader": {
     "grade": false,
     "grade_id": "cell-2f80ed5f342a854e",
     "locked": true,
     "schema_version": 3,
     "solution": false,
     "task": false
    }
   },
   "outputs": [],
   "source": [
    "# plot average and standard deviation of 50 runs\n",
    "repeats = len(train_logs_ppo_custom['envs'])\n",
    "rewards_mat = np.asarray([_env.get_episode_rewards()[:250] for _env in train_logs_ppo_custom['envs']])\n",
    "rewards_average = np.median(rewards_mat, axis=0)\n",
    "rewards_std = rewards_mat.std(axis=0)\n",
    "\n",
    "plt.figure()\n",
    "ax = plt.gca()\n",
    "ax.plot(rewards_average, color='green', label='median reward')\n",
    "ax.plot(rewards_mat.max(axis=0), color='cyan', label='best seen reward')\n",
    "ax.set_ylabel('Reward')\n",
    "ax.set_xlabel('Episode')\n",
    "ax.set_ylim(0, 0.015)\n",
    "ax.axhline(OPTIMAL_CONTROL, ls='--', color='red', label='optimal')\n",
    "ax.axhline(RANDOM_BEST, ls='--', color='orange', label='best rnd actions')\n",
    "ax.axhline(RANDOM_AVG, ls='--', color='yellow', label='average rnd actions')\n",
    "ax.legend(loc='lower left')\n",
    "ax.set_title('Summary of 50 experiments')\n",
    "plt.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d7b37741",
   "metadata": {
    "nbgrader": {
     "grade": false,
     "grade_id": "cell-fe61925f62283573",
     "locked": true,
     "schema_version": 3,
     "solution": false,
     "task": false
    }
   },
   "source": [
    "Very strong performance is sometimes achieved even within 10 training episodes. Here, the 'optimal' performance is based on an analytical solution using a continuous-time model (planning) which is applied to an approximate discrete-time rocket simulation. Hence, there is a systematic deviation between the model of the analytical solution and the actual simulation behavior (i.e., it is actually non-optimal w.r.t. the simplified simulation).\n",
    "\n",
    "The low median reward curve shows that several repetitions of an algorithm are inevitable for a robust assessment."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c0cf4a7e",
   "metadata": {
    "nbgrader": {
     "grade": false,
     "grade_id": "cell-cd9a6dc2f628c73d",
     "locked": true,
     "schema_version": 3,
     "solution": false,
     "task": false
    }
   },
   "source": [
    "### Demo: StableBaselines3 usage\n",
    "\n",
    "Alternatively you can use readily available Python packages such as [stable-baselines3](https://github.com/DLR-RM/stable-baselines3) for employment of state-of-the-art algorithms.\n",
    "In what follows below, the PPO algorithm as utilized by stable-baselines3 is showcased."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1b5c9fb7",
   "metadata": {
    "nbgrader": {
     "grade": false,
     "grade_id": "cell-b91e000d33246cbb",
     "locked": true,
     "schema_version": 3,
     "solution": false,
     "task": false
    }
   },
   "outputs": [],
   "source": [
    "class Pbar(BaseCallback):\n",
    "    def __init__(self, total, n_rollout_steps):\n",
    "        super().__init__()\n",
    "        self.n_steps = n_rollout_steps\n",
    "        self.pbar = tqdm(desc='Training', total=total_timesteps, position=0, leave=True)\n",
    "\n",
    "    def _on_rollout_end(self) -> None:\n",
    "        \"\"\"This event is triggered before updating the policy.\"\"\"\n",
    "        pass\n",
    "\n",
    "    def _on_step(self) -> bool:\n",
    "        self.pbar.update()\n",
    "        return True\n",
    "    \n",
    "    def _on_training_end(self) -> None:\n",
    "        \"\"\"This event is triggered before exiting the `learn()` method.\"\"\"\n",
    "        print(\"\")\n",
    "\n",
    "    \n",
    "policy_kwargs = dict(activation_fn=torch.nn.Sigmoid, net_arch=[dict(pi=[8], vf=[8])])\n",
    "\n",
    "n_steps = 256  # an episode takes around ~250 steps\n",
    "total_timesteps = n_steps*250  # steps_per_episode * episodes\n",
    "\n",
    "repeats = 4  # repeat experiment x times to assess scatter through random init\n",
    "train_logs_ppo_sb3 = {'envs': [], 'models': []}\n",
    "for _ in range(repeats):\n",
    "    \n",
    "    env = Monitor(GoddardEnv(render_mode = None))\n",
    "    model = PPO('MlpPolicy', env, n_steps=n_steps, n_epochs=8,  verbose=0, device='cpu', learning_rate=1e-3, batch_size=n_steps, policy_kwargs=policy_kwargs,\n",
    "               clip_range=0.1, clip_range_vf=0.1, gamma=0.999)\n",
    "    model = model.learn(total_timesteps=total_timesteps, callback=Pbar(total_timesteps, n_steps))\n",
    "    train_logs_ppo_sb3['envs'] += [env]\n",
    "    train_logs_ppo_sb3['models'] += [model]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "96caae1a",
   "metadata": {
    "nbgrader": {
     "grade": false,
     "grade_id": "cell-531e062c849ee81b",
     "locked": true,
     "schema_version": 3,
     "solution": false,
     "task": false
    }
   },
   "outputs": [],
   "source": [
    "plot_reward_trends(train_logs_ppo_sb3)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cf1a303a",
   "metadata": {
    "nbgrader": {
     "grade": false,
     "grade_id": "cell-78eedb7acc89df22",
     "locked": true,
     "schema_version": 3,
     "solution": false,
     "task": false
    }
   },
   "source": [
    "As you can see, the chosen architecture and hyper parameters do well often, but can also fail abruptly. \n",
    "What might be the reason for these failures?"
   ]
  }
 ],
 "metadata": {
  "celltoolbar": "Create Assignment",
  "kernelspec": {
   "display_name": "RLKurs2",
   "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.9.19"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}