Reinforcement_Learning/Proximal_Policy_Optimization/Proximal_Policy_Optimization.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "responsible-likelihood",
   "metadata": {},
   "source": [
    "# Proximal Policy Optimization"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ordinary-ballet",
   "metadata": {},
   "source": [
    "## Introduction"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "active-stress",
   "metadata": {},
   "source": [
    "### Policy gradient"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "stone-sheffield",
   "metadata": {},
   "source": [
    "$$J(\\theta)=\\mathbb{E}_{\\pi_\\theta} \\left [ \\sum_t r_t \\right ] =v_{\\pi_\\theta}(s_0)=\\sum_{s\\in S}d(s)*v_{\\pi_\\theta}(s)$$"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "separated-consequence",
   "metadata": {},
   "source": [
    "1step MDP : \n",
    "\n",
    "$$J(\\theta)=\\sum_{s\\in S}d(s)*v_{\\pi_\\theta}(s)=\\sum_{s\\in S}d(s)\\sum_{a\\in A}\\pi_\\theta(s,a)*R_{s,a}$$\n",
    "\n",
    "$$\\begin{align*}\n",
    "\\triangledown_\\theta J(\\theta) &= \\triangledown_\\theta \\sum_{s\\in S}d(s)\\sum_{a\\in A}\\pi_\\theta(s,a)*R_{s,a} \\\\\n",
    "&= \\sum_{s\\in S}d(s)\\sum_{a\\in A}\\triangledown_\\theta\\pi_\\theta(s,a)*R_{s,a} \\\\\n",
    "&=\\sum_{s\\in S}d(s)\\sum_{a\\in A}\\frac{\\pi_\\theta (s,a)}{\\pi_\\theta (s,a)}\\triangledown_\\theta\\pi_\\theta(s,a)*R_{s,a} \\\\\n",
    "&=\\sum_{s\\in S}d(s)\\sum_{a\\in A}\\pi_\\theta (s,a)\\frac{\\triangledown_\\theta\\pi_\\theta(s,a)}{\\pi_\\theta (s,a)}*R_{s,a} \\\\\n",
    "&=\\sum_{s\\in S}d(s)\\sum_{a\\in A}\\pi_\\theta (s,a)\\triangledown_\\theta log\\pi_\\theta(s,a)*R_{s,a} \\\\\n",
    "&=\\mathbb{E}_{\\pi_\\theta}[\\triangledown_\\theta log\\pi_\\theta(s,a)*R_{s,a}]\n",
    "\\end{align*}$$"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "certain-viewer",
   "metadata": {},
   "source": [
    "MDP : $$\\triangledown_\\theta J(\\theta) = \\mathbb{E}_{\\pi_\\theta}[\\triangledown_\\theta log\\pi_\\theta(s,a)*Q_{\\pi_\\theta}(s,a)]$$"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "protecting-invention",
   "metadata": {},
   "source": [
    "### REINFOCE  Algorithm"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "biological-elements",
   "metadata": {},
   "source": [
    "$G_t$의 샘플을 여러 개 얻어서 평균을 내면 그 값이 실제 $Q_{\\pi_\\theta}(s,a)$에 근사하기 때문"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "confidential-intranet",
   "metadata": {},
   "source": [
    "$$\\triangledown_\\theta J(\\theta) = \\mathbb{E}_{\\pi_\\theta}[\\triangledown_\\theta log\\pi_\\theta(s,a)*G_t]$$"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "painted-northeast",
   "metadata": {},
   "source": [
    "### Advantage Policy gradient"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "german-puzzle",
   "metadata": {},
   "source": [
    "$$\\triangledown_\\theta J(\\theta) = \\mathbb{E}_{\\pi_\\theta}[\\triangledown_\\theta log\\pi_\\theta(s,a)*\\{Q_{\\pi_\\theta}(s,a)-V_{\\pi_\\theta}(s)\\}]$$\n",
    "$$\\triangledown_\\theta J(\\theta) = \\mathbb{E}_{\\pi_\\theta}[\\triangledown_\\theta log\\pi_\\theta(s,a)*A_{\\pi_\\theta}(s,a)]$$"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "artistic-orbit",
   "metadata": {},
   "source": [
    "### Advantage Policy gradient loss"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "innocent-lodge",
   "metadata": {},
   "source": [
    "$$L^{PG}(\\theta)=\\mathbb{E}_{Q_{\\pi_\\theta}}[log\\pi_\\theta(s,a)*A_{\\pi_\\theta}(s,a)]$$"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "related-museum",
   "metadata": {},
   "source": [
    "### Importance Sampling"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "charged-flood",
   "metadata": {},
   "source": [
    "$$\\triangledown_\\theta log\\pi_\\theta(s,a)|_{\\theta_\\text{old}}=\\frac{\\triangledown_\\theta\\pi_\\theta(s,a)|_{\\theta_\\text{old}}}{\\pi_{\\theta_\\text{old}}(s,a)}=\\triangledown_{\\theta}\\begin{pmatrix}\\frac{\\pi_\\theta(s,a)}{\\pi_{\\theta_\\text{old}}(s,a)}\\end{pmatrix}|_{\\theta_\\text{old}}$$\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "organized-night",
   "metadata": {},
   "source": [
    "$$L^{IS}_{\\theta_{\\text{old}}}(\\theta)=\\mathbb{E}_{Q_{\\pi_{\\theta_\\text{old}}}}\\begin{bmatrix}\\frac{\\pi_\\theta(s,a)}{\\pi_{\\theta_\\text{old}}(s,a)}*A_{\\pi_\\theta}(s,a)\\end{bmatrix}$$"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "answering-biology",
   "metadata": {},
   "source": [
    "### Cliping"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "pleasant-manual",
   "metadata": {},
   "source": [
    "$\\theta$와 $\\theta_{\\text{old}}$는 비슷해야한다."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "brown-softball",
   "metadata": {},
   "source": [
    "$$L^{CLIP}(\\theta) = \\mathbb{E}_t[\\text{min}(r_t(\\theta)A_t, \\text{clip}(r_t(\\theta),1-\\epsilon,1+\\epsilon)A_t)]$$"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "frequent-mechanics",
   "metadata": {},
   "source": [
    "<img src='../img/PPO_1.png' width='800'> "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "yellow-diesel",
   "metadata": {},
   "source": [
    "### Adaptive KL Penalty Coefficient"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "respiratory-tattoo",
   "metadata": {},
   "source": [
    "<img src='../img/ppo_2.png' width='800'> "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "nuclear-postcard",
   "metadata": {},
   "source": [
    "### Generalized Advantage Estimation"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "whole-lunch",
   "metadata": {},
   "source": [
    "n-step TD 대신 GAE를 사용"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "medieval-attitude",
   "metadata": {},
   "source": [
    "$$\\begin{matrix}\n",
    "R_t+\\gamma V(s_{t+1})-V(s_t) \\text{  : 1-step TD error} \\\\\n",
    "R_t+\\gamma R_{t+1}+\\gamma^2V(s_{t+2})-V(s_t) \\text{  : 2-step TD error} \\\\\n",
    "\\vdots \\\\\n",
    "\\end{matrix}$$"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ethical-thursday",
   "metadata": {},
   "source": [
    "1개만 이용하는게 아니라 exponential moving average 하자."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "compact-thousand",
   "metadata": {},
   "source": [
    "간단하게 나타내기 위해 $A_t^{(1)}=\\delta_t$로 치환하자 "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "lesbian-southwest",
   "metadata": {},
   "source": [
    "$$A_t^{(2)}=R_t+\\gamma R_{t+1}+\\gamma^2V(s_{t+2})-V(s_t)=(R_t+\\gamma V(s_{t+1})-V(s_t))+\\gamma(R_{t+1}+\\gamma V(s_{t+2})-V(s_{t+1}))=\\delta_t+\\gamma\\delta_{t+1}$$"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "pharmaceutical-report",
   "metadata": {},
   "source": [
    "$$A_t^{(n)}=\\sum_{k=t}^{t+n-1}\\gamma^{k-t}\\delta_k$$ 로 표현할 수 있다."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "hindu-collector",
   "metadata": {},
   "source": [
    "$$TD(\\lambda)\\rightarrow\\sum^\\infty_{n=1}(1-\\lambda)\\lambda^{n-1}A_t^{(n)}$$를 이와 같이 표현하면"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "radio-duplicate",
   "metadata": {},
   "source": [
    "$$\\begin{matrix}\n",
    "TD(\\lambda)&=(1-\\lambda)(\\delta_t+\\lambda(\\delta_t+\\gamma\\delta_{t+1})+\\lambda^2(\\delta_t+\\gamma\\delta_{t+1}+\\gamma^2\\delta_{t+2})+\\cdots) \\\\\n",
    "&=(1-\\lambda)(\\delta_t(1+\\lambda+\\lambda^2+\\cdots)+\\gamma\\delta_{t+1}(\\lambda+\\lambda^2+\\cdots)+\\cdots) \\\\\n",
    "&=(1-\\lambda)(\\delta_t*\\frac{1}{1-\\lambda}+\\gamma\\delta_{t+1}*\\frac{\\lambda}{1-\\lambda}+\\cdots) \\\\ \n",
    "&=\\sum_{k=t}^{\\infty}(\\gamma\\lambda)^{k-t}\\delta_k\n",
    "\\end{matrix}$$로 표현될 수 있다."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "prompt-pressing",
   "metadata": {},
   "source": [
    "하지만 몬테카를로가 아니라 n-step TD이기 때문에 $\\infty$를 계산할 수 없다. 따라서 근사를 사용"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "secure-fight",
   "metadata": {},
   "source": [
    "$$\\begin{matrix}\n",
    "A_i^{GAE}=\\sum_{k=i}^{\\infty}(\\gamma\\lambda)^{k-i}\\delta_k \\\\\n",
    "\\hat{A}_i^{GAE}=\\sum_{k=i}^{t}(\\gamma\\lambda)^{k-i}\\delta_k\n",
    "\\end{matrix}$$"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "enormous-shadow",
   "metadata": {},
   "source": [
    "<img src='https://i.imgur.com/rd5tda1.png' width='800'> "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "optical-romantic",
   "metadata": {},
   "source": [
    "<img src='../img/ppo.png' width='800'> "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "enabling-office",
   "metadata": {},
   "source": [
    "## PPO"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "adaptive-surfing",
   "metadata": {},
   "source": [
    "### Libraries"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "tracked-duration",
   "metadata": {},
   "source": [
    "For this example the following libraries are used:\n",
    "\n",
    "1. numpy for n-dimensional arrays\n",
    "2. tensorflow and keras for building the deep RL PPO agent\n",
    "3. gym for getting everything we need about the environment\n",
    "4. scipy.signal for calculating the discounted cumulative sums of vectors"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "id": "attached-accordance",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-07-24T20:17:30.200329Z",
     "start_time": "2021-07-24T20:17:30.195684Z"
    }
   },
   "outputs": [],
   "source": [
    "import gym\n",
    "import numpy as np\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "from torch.optim import Adam\n",
    "from torch.distributions import MultivariateNormal\n",
    "from torch.distributions import Categorical\n",
    "import scipy.signal\n",
    "from mpi4py import MPI\n",
    "import time\n",
    "import os\n",
    "import sys\n",
    "import subprocess\n",
    "from gym.spaces import Box, Discrete"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "impressed-memorabilia",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-07-24T19:37:32.020285Z",
     "start_time": "2021-07-24T19:37:32.017140Z"
    }
   },
   "outputs": [],
   "source": [
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "neutral-integrity",
   "metadata": {},
   "source": [
    "### Functions and class"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "accessible-january",
   "metadata": {},
   "source": [
    "- Discount Factor 적용"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "recent-gazette",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-07-24T19:37:32.769897Z",
     "start_time": "2021-07-24T19:37:32.763151Z"
    }
   },
   "outputs": [],
   "source": [
    "def discounted_cumulative_sums(x, discount):\n",
    "    \"\"\"\n",
    "    magic from rllab for computing discounted cumulative sums of vectors.\n",
    "    input: \n",
    "        vector x, \n",
    "        [x0, \n",
    "         x1, \n",
    "         x2]\n",
    "    output:\n",
    "        [x0 + discount * x1 + discount^2 * x2,  \n",
    "         x1 + discount * x2,\n",
    "         x2]\n",
    "    \"\"\"\n",
    "    return scipy.signal.lfilter([1], [1, float(-discount)], x[::-1], axis=0)[::-1]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "czech-pillow",
   "metadata": {},
   "source": [
    "```pyhon\n",
    "scipy.signal.lfilter(b, a, x, axis=- 1, zi=None)\n",
    "```\n",
    "\n",
    "$$a[0]*y[n] = b[0]*x[n] + b[1]*x[n-1] + ... + b[M]*x[n-M]\n",
    "                      - a[1]*y[n-1] - ... - a[N]*y[n-N]$$"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "simplified-mailing",
   "metadata": {},
   "source": [
    "- step size X dim shape 생성 function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 54,
   "id": "hidden-victorian",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-07-24T21:08:39.668676Z",
     "start_time": "2021-07-24T21:08:39.664014Z"
    }
   },
   "outputs": [],
   "source": [
    "def combined_shape(length, shape=None):\n",
    "    if shape is None:\n",
    "        return (length,)\n",
    "    return (length, shape) if np.isscalar(shape) else (length, *shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "romantic-possibility",
   "metadata": {},
   "source": [
    "- Buffer 생성, 저장, 추출 기능을 제공하는 class"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 57,
   "id": "progressive-gambling",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-07-24T21:09:29.259807Z",
     "start_time": "2021-07-24T21:09:29.241464Z"
    }
   },
   "outputs": [],
   "source": [
    "class PPOBuffer:\n",
    "    # Buffer for storing trajectories\n",
    "    def __init__(self, obs_dim, act_dim, size, gamma=0.99, lam=0.95):\n",
    "        # Buffer initialization\n",
    "        self.observation_buffer = np.zeros(combined_shape(size, obs_dim), dtype=np.float32)\n",
    "        self.action_buffer = np.zeros(combined_shape(size, act_dim), dtype=np.float32)\n",
    "        self.advantage_buffer = np.zeros(size, dtype=np.float32)\n",
    "        self.reward_buffer = np.zeros(size, dtype=np.float32)\n",
    "        self.return_buffer = np.zeros(size, dtype=np.float32)\n",
    "        self.value_buffer = np.zeros(size, dtype=np.float32)\n",
    "        self.logprobability_buffer = np.zeros(size, dtype=np.float32)\n",
    "        self.gamma, self.lam = gamma, lam\n",
    "        self.pointer, self.trajectory_start_index, self.max_size = 0, 0, size\n",
    "\n",
    "    def store(self, observation, action, reward, value, logprobability):\n",
    "        # Append one step of agent-environment interaction\n",
    "        self.observation_buffer[self.pointer] = observation\n",
    "        self.action_buffer[self.pointer] = action\n",
    "        self.reward_buffer[self.pointer] = reward\n",
    "        self.value_buffer[self.pointer] = value\n",
    "        self.logprobability_buffer[self.pointer] = logprobability\n",
    "        self.pointer += 1\n",
    "\n",
    "    def finish_trajectory(self, last_value=0):\n",
    "        # Finish the trajectory by computing advantage estimates and rewards-to-go\n",
    "        path_slice = slice(self.trajectory_start_index, self.pointer)\n",
    "        rewards = np.append(self.reward_buffer[path_slice], last_value)\n",
    "        values = np.append(self.value_buffer[path_slice], last_value)\n",
    "\n",
    "        deltas = rewards[:-1] + self.gamma * values[1:] - values[:-1]\n",
    "\n",
    "        self.advantage_buffer[path_slice] = discounted_cumulative_sums(\n",
    "            deltas, self.gamma * self.lam\n",
    "        )\n",
    "        self.return_buffer[path_slice] = discounted_cumulative_sums(\n",
    "            rewards, self.gamma\n",
    "        )[:-1]\n",
    "\n",
    "        self.trajectory_start_index = self.pointer\n",
    "    \n",
    "    def get(self):\n",
    "        # Get all data of the buffer and normalize the advantages\n",
    "        assert self.pointer == self.max_size\n",
    "        self.pointer, self.trajectory_start_index = 0, 0\n",
    "        advantage_mean, advantage_std = (\n",
    "            np.mean(self.advantage_buffer),\n",
    "            np.std(self.advantage_buffer),\n",
    "        )\n",
    "        self.advantage_buffer = (self.advantage_buffer - advantage_mean) / advantage_std\n",
    "        data = dict(obs=self.observation_buffer, \n",
    "                    act=self.action_buffer, \n",
    "                    ret=self.return_buffer,\n",
    "                    adv=self.advantage_buffer, \n",
    "                    logp=self.logprobability_buffer)\n",
    "        return {k: torch.as_tensor(v, dtype=torch.float32) for k,v in data.items()}\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "abstract-question",
   "metadata": {},
   "source": [
    "- 모델 네트워크 선언"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "rocky-context",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-07-24T19:37:34.321413Z",
     "start_time": "2021-07-24T19:37:34.276113Z"
    }
   },
   "outputs": [],
   "source": [
    "def mlp(sizes, activation, output_activation=nn.Identity):\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",
    "\n",
    "class Actor(nn.Module):\n",
    "\n",
    "    def _distribution(self, obs):\n",
    "        raise NotImplementedError\n",
    "\n",
    "    def _log_prob_from_distribution(self, pi, act):\n",
    "        raise NotImplementedError\n",
    "\n",
    "    def forward(self, obs, act=None):\n",
    "        # Produce action distributions for given observations, and \n",
    "        # optionally compute the log likelihood of given actions under\n",
    "        # those distributions.\n",
    "        pi = self._distribution(obs)\n",
    "        logp_a = None\n",
    "        if act is not None:\n",
    "            logp_a = self._log_prob_from_distribution(pi, act)\n",
    "        return pi, logp_a\n",
    "\n",
    "\n",
    "class MLPCategoricalActor(Actor):\n",
    "    \n",
    "    def __init__(self, obs_dim, act_dim, hidden_sizes, activation):\n",
    "        super().__init__()\n",
    "        self.logits_net = mlp([obs_dim] + list(hidden_sizes) + [act_dim], activation)\n",
    "\n",
    "    def _distribution(self, obs):\n",
    "        logits = self.logits_net(obs)\n",
    "        return Categorical(logits=logits)\n",
    "\n",
    "    def _log_prob_from_distribution(self, pi, act):\n",
    "        return pi.log_prob(act)\n",
    "\n",
    "\n",
    "class MLPGaussianActor(Actor):\n",
    "\n",
    "    def __init__(self, obs_dim, act_dim, hidden_sizes, activation):\n",
    "        super().__init__()\n",
    "        log_std = -0.5 * np.ones(act_dim, dtype=np.float32)\n",
    "        self.log_std = torch.nn.Parameter(torch.as_tensor(log_std))\n",
    "        self.mu_net = mlp([obs_dim] + list(hidden_sizes) + [act_dim], activation)\n",
    "\n",
    "    def _distribution(self, obs):\n",
    "        mu = self.mu_net(obs)\n",
    "        std = torch.exp(self.log_std)\n",
    "        return Normal(mu, std)\n",
    "\n",
    "    def _log_prob_from_distribution(self, pi, act):\n",
    "        return pi.log_prob(act).sum(axis=-1)    # Last axis sum needed for Torch Normal distribution\n",
    "    \n",
    "    \n",
    "class MLPCritic(nn.Module):\n",
    "\n",
    "    def __init__(self, obs_dim, hidden_sizes, activation):\n",
    "        super().__init__()\n",
    "        self.v_net = mlp([obs_dim] + list(hidden_sizes) + [1], activation)\n",
    "\n",
    "    def forward(self, obs):\n",
    "        return torch.squeeze(self.v_net(obs), -1) # Critical to ensure v has right shape."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "ancient-doctrine",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-07-24T19:37:35.056102Z",
     "start_time": "2021-07-24T19:37:35.048510Z"
    }
   },
   "outputs": [],
   "source": [
    "class MLPActorCritic(nn.Module):\n",
    "\n",
    "\n",
    "    def __init__(self, observation_space, action_space, \n",
    "                 hidden_sizes=(64,64), activation=nn.Tanh):\n",
    "        super().__init__()\n",
    "\n",
    "        obs_dim = observation_space.shape[0]\n",
    "\n",
    "        # policy builder depends on action space\n",
    "        if isinstance(action_space, Box):\n",
    "            self.pi = MLPGaussianActor(obs_dim, action_space.shape[0], hidden_sizes, activation)\n",
    "        elif isinstance(action_space, Discrete):\n",
    "            self.pi = MLPCategoricalActor(obs_dim, action_space.n, hidden_sizes, activation)\n",
    "\n",
    "        # build value function\n",
    "        self.v  = MLPCritic(obs_dim, hidden_sizes, activation)\n",
    "\n",
    "    def step(self, obs):\n",
    "        with torch.no_grad():\n",
    "            pi = self.pi._distribution(obs)\n",
    "            a = pi.sample()\n",
    "            logp_a = self.pi._log_prob_from_distribution(pi, a)\n",
    "            v = self.v(obs)\n",
    "        return a.numpy(), v.numpy(), logp_a.numpy()\n",
    "\n",
    "    def act(self, obs):\n",
    "        return self.step(obs)[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 89,
   "id": "rental-porcelain",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-07-25T18:00:59.391662Z",
     "start_time": "2021-07-25T18:00:59.311240Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'../img/ppo_v.png'"
      ]
     },
     "execution_count": 89,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from torchviz import make_dot\n",
    "from torch.autograd import Variable\n",
    "\n",
    "ac = MLPActorCritic(gym.make(\"CartPole-v0\").observation_space, \n",
    "                    gym.make(\"CartPole-v0\").action_space, \n",
    "                    hidden_sizes=hidden_sizes)\n",
    "act=0\n",
    "env = gym.make(\"CartPole-v0\")\n",
    "env.reset()\n",
    "obs, _, _, _ = env.step(act)\n",
    "# make_dot(ac.pi(torch.as_tensor(obs, dtype=torch.float32), torch.as_tensor(act)), \n",
    "#          params=dict(ac.pi.named_parameters())).render(\"../img/ppo_pi\", format=\"png\")\n",
    "make_dot(ac.v(torch.as_tensor(obs, dtype=torch.float32)), \n",
    "         params=dict(ac.v.named_parameters())).render(\"../img/ppo_v\", format=\"png\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "grand-publication",
   "metadata": {},
   "source": [
    "<img src='../img/ppo_v.png' width='400'> "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "narrative-diagnosis",
   "metadata": {},
   "source": [
    "- 분산 및 병렬처리를 위한 MPI 사용 function 선언\n",
    "[mpi4py example](https://fortran.readthedocs.io/ko/latest/mpi4py_ex/)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 63,
   "id": "unusual-speed",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-07-24T21:12:38.442308Z",
     "start_time": "2021-07-24T21:12:38.428832Z"
    }
   },
   "outputs": [],
   "source": [
    "def num_procs():\n",
    "    \"\"\"Count active MPI processes.\"\"\"\n",
    "    return MPI.COMM_WORLD.Get_size()\n",
    "\n",
    "def setup_pytorch_for_mpi():\n",
    "    \"\"\"\n",
    "    Avoid slowdowns caused by each separate process's PyTorch using\n",
    "    more than its fair share of CPU resources.\n",
    "    \"\"\"\n",
    "    #print('Proc %d: Reporting original number of Torch threads as %d.'%(proc_id(), torch.get_num_threads()), flush=True)\n",
    "    if torch.get_num_threads()==1:\n",
    "        return\n",
    "    fair_num_threads = max(int(torch.get_num_threads() / num_procs()), 1)\n",
    "    torch.set_num_threads(fair_num_threads)\n",
    "    #print('Proc %d: Reporting new number of Torch threads as %d.'%(proc_id(), torch.get_num_threads()), flush=True)\n",
    "    \n",
    "def sync_params(module):\n",
    "    \"\"\" Sync all parameters of module across all MPI processes. \"\"\"\n",
    "    if num_procs()==1:\n",
    "        return\n",
    "    for p in module.parameters():\n",
    "        p_numpy = p.data.numpy()\n",
    "        broadcast(p_numpy)\n",
    "        \n",
    "def mpi_avg_grads(module):\n",
    "    \"\"\" Average contents of gradient buffers across MPI processes. \"\"\"\n",
    "    if num_procs()==1:\n",
    "        return\n",
    "    for p in module.parameters():\n",
    "        p_grad_numpy = p.grad.numpy()   # numpy view of tensor data\n",
    "        avg_p_grad = mpi_avg(p.grad)\n",
    "        p_grad_numpy[:] = avg_p_grad[:]\n",
    "        \n",
    "def mpi_fork(n, bind_to_core=False):\n",
    "    \"\"\"\n",
    "    Re-launches the current script with workers linked by MPI.\n",
    "    Also, terminates the original process that launched it.\n",
    "    Taken almost without modification from the Baselines function of the\n",
    "    `same name`_.\n",
    "    .. _`same name`: https://github.com/openai/baselines/blob/master/baselines/common/mpi_fork.py\n",
    "    Args:\n",
    "        n (int): Number of process to split into.\n",
    "        bind_to_core (bool): Bind each MPI process to a core.\n",
    "    \"\"\"\n",
    "    if n<=1: \n",
    "        return\n",
    "    if os.getenv(\"IN_MPI\") is None:\n",
    "        env = os.environ.copy()\n",
    "        env.update(\n",
    "            MKL_NUM_THREADS=\"1\",\n",
    "            OMP_NUM_THREADS=\"1\",\n",
    "            IN_MPI=\"1\"\n",
    "        )\n",
    "        args = [\"mpirun\", \"-np\", str(n)]\n",
    "        if bind_to_core:\n",
    "            args += [\"-bind-to\", \"core\"]\n",
    "        args += [sys.executable] + sys.argv\n",
    "        subprocess.check_call(args, env=env)\n",
    "        sys.exit()\n",
    "        \n",
    "def allreduce(*args, **kwargs):\n",
    "    return MPI.COMM_WORLD.Allreduce(*args, **kwargs)\n",
    "    \n",
    "def mpi_op(x, op):\n",
    "    x, scalar = ([x], True) if np.isscalar(x) else (x, False)\n",
    "    x = np.asarray(x, dtype=np.float32)\n",
    "    buff = np.zeros_like(x, dtype=np.float32)\n",
    "    allreduce(x, buff, op=op)\n",
    "    return buff[0] if scalar else buff\n",
    "        \n",
    "def mpi_sum(x):\n",
    "    return mpi_op(x, MPI.SUM)        \n",
    "        \n",
    "def mpi_avg(x):\n",
    "    \"\"\"Average a scalar or vector over MPI processes.\"\"\"\n",
    "    return mpi_sum(x) / num_procs()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "informational-irrigation",
   "metadata": {},
   "source": [
    "- PPO training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 67,
   "id": "atomic-arabic",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-07-25T01:14:11.273226Z",
     "start_time": "2021-07-25T01:14:11.248778Z"
    }
   },
   "outputs": [],
   "source": [
    "def ppo(env_fn, actor_critic=MLPActorCritic, ac_kwargs=dict(), seed=702, \n",
    "        steps_per_epoch=4000, epochs=50, gamma=0.99, clip_ratio=0.2, pi_lr=3e-4,\n",
    "        vf_lr=1e-3, train_pi_iters=80, train_v_iters=80, lam=0.97, max_ep_len=1000,\n",
    "        target_kl=0.01):\n",
    "    # Special function to avoid certain slowdowns from PyTorch + MPI combo.\n",
    "    setup_pytorch_for_mpi()\n",
    "\n",
    "    # Random seed\n",
    "    seed = 702\n",
    "    torch.manual_seed(seed)\n",
    "    np.random.seed(seed)\n",
    "\n",
    "    # Instantiate environment\n",
    "    env = env_fn()\n",
    "    obs_dim = env.observation_space.shape[0]\n",
    "    act_dim = env.action_space.shape\n",
    "\n",
    "    # Create actor-critic module\n",
    "    ac = actor_critic(env.observation_space, env.action_space, **ac_kwargs)\n",
    "\n",
    "    # Sync params across processes\n",
    "    sync_params(ac)\n",
    "\n",
    "    # Count variables\n",
    "    var_counts = tuple(sum([np.prod(p.shape) for p in module.parameters()]) for module in [ac.pi, ac.v])\n",
    "\n",
    "    # Set up experience buffer\n",
    "    local_steps_per_epoch = int(steps_per_epoch / num_procs())\n",
    "    buf = PPOBuffer(obs_dim, act_dim, local_steps_per_epoch, gamma, lam)\n",
    "\n",
    "    # Set up function for computing PPO policy loss\n",
    "    def compute_loss_pi(data):\n",
    "        obs, act, adv, logp_old = data['obs'], data['act'], data['adv'], data['logp']\n",
    "\n",
    "        # Policy loss\n",
    "        pi, logp = ac.pi(obs, act)\n",
    "        ratio = torch.exp(logp - logp_old)\n",
    "        clip_adv = torch.clamp(ratio, 1-clip_ratio, 1+clip_ratio) * adv\n",
    "        loss_pi = -(torch.min(ratio * adv, clip_adv)).mean()\n",
    "\n",
    "        # Useful extra info\n",
    "        approx_kl = (logp_old - logp).mean().item()\n",
    "        ent = pi.entropy().mean().item()\n",
    "        clipped = ratio.gt(1+clip_ratio) | ratio.lt(1-clip_ratio)\n",
    "        clipfrac = torch.as_tensor(clipped, dtype=torch.float32).mean().item()\n",
    "        pi_info = dict(kl=approx_kl, ent=ent, cf=clipfrac)\n",
    "\n",
    "        return loss_pi, pi_info\n",
    "\n",
    "    # Set up function for computing value loss\n",
    "    def compute_loss_v(data):\n",
    "        obs, ret = data['obs'], data['ret']\n",
    "        return ((ac.v(obs) - ret)**2).mean()\n",
    "\n",
    "    # Set up optimizers for policy and value function\n",
    "    pi_optimizer = Adam(ac.pi.parameters(), lr=pi_lr)\n",
    "    vf_optimizer = Adam(ac.v.parameters(), lr=vf_lr)\n",
    "\n",
    "    \n",
    "    def update():\n",
    "        data = buf.get()\n",
    "\n",
    "        pi_l_old, pi_info_old = compute_loss_pi(data)\n",
    "        pi_l_old = pi_l_old.item()\n",
    "        v_l_old = compute_loss_v(data).item()\n",
    "\n",
    "        # Train policy with multiple steps of gradient descent\n",
    "        for i in range(train_pi_iters):\n",
    "            pi_optimizer.zero_grad()\n",
    "            loss_pi, pi_info = compute_loss_pi(data)\n",
    "            kl = mpi_avg(pi_info['kl'])\n",
    "            if kl > 1.5 * target_kl:\n",
    "                break\n",
    "            loss_pi.backward()\n",
    "            mpi_avg_grads(ac.pi)    # average grads across MPI processes\n",
    "            pi_optimizer.step()\n",
    "\n",
    "\n",
    "        # Value function learning\n",
    "        for i in range(train_v_iters):\n",
    "            vf_optimizer.zero_grad()\n",
    "            loss_v = compute_loss_v(data)\n",
    "            loss_v.backward()\n",
    "            mpi_avg_grads(ac.v)    # average grads across MPI processes\n",
    "            vf_optimizer.step()\n",
    "\n",
    "\n",
    "    # Prepare for interaction with environment\n",
    "    start_time = time.time()\n",
    "    o, ep_ret, ep_len = env.reset(), 0, 0\n",
    "\n",
    "    # Main loop: collect experience in env and update/log each epoch\n",
    "    for epoch in range(epochs):\n",
    "        for t in range(local_steps_per_epoch):\n",
    "            a, v, logp = ac.step(torch.as_tensor(o, dtype=torch.float32))\n",
    "\n",
    "            next_o, r, d, _ = env.step(a)\n",
    "            ep_ret += r\n",
    "            ep_len += 1\n",
    "\n",
    "            # save and log\n",
    "            buf.store(o, a, r, v, logp)\n",
    "            \n",
    "            # Update obs (critical!)\n",
    "            o = next_o\n",
    "\n",
    "            timeout = ep_len == max_ep_len\n",
    "            terminal = d or timeout\n",
    "            epoch_ended = t==local_steps_per_epoch-1\n",
    "\n",
    "            if terminal or epoch_ended:\n",
    "                if epoch_ended and not(terminal):\n",
    "                    print(f'Epoch: {epoch} Warning: trajectory cut off by epoch at {ep_len} steps.', flush=True)\n",
    "                # if trajectory didn't reach terminal state, bootstrap value target\n",
    "                if timeout or epoch_ended:\n",
    "                    _, v, _ = ac.step(torch.as_tensor(o, dtype=torch.float32))\n",
    "                else:\n",
    "                    v = 0\n",
    "                buf.finish_trajectory(v)\n",
    "                o, ep_ret, ep_len = env.reset(), 0, 0\n",
    "\n",
    "\n",
    "        # Perform PPO update!\n",
    "        update()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "afraid-pulse",
   "metadata": {},
   "source": [
    "### Hyperparameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 65,
   "id": "polish-frank",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-07-25T01:13:01.518399Z",
     "start_time": "2021-07-25T01:13:01.514619Z"
    }
   },
   "outputs": [],
   "source": [
    "# Hyperparameters of the PPO algorithm\n",
    "steps_per_epoch = 4000\n",
    "epochs = 100\n",
    "gamma = 0.99\n",
    "clip_ratio = 0.2\n",
    "policy_learning_rate = 3e-4\n",
    "value_function_learning_rate = 1e-3\n",
    "train_policy_iterations = 80\n",
    "train_value_iterations = 80\n",
    "lam = 0.97\n",
    "target_kl = 0.01\n",
    "hidden_sizes = (64, 64)\n",
    "num_cpu = 1\n",
    "# True if you want to render the environment\n",
    "render = False"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "crucial-listening",
   "metadata": {},
   "source": [
    "### Training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 68,
   "id": "moving-cement",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-07-25T01:20:43.470503Z",
     "start_time": "2021-07-25T01:14:14.105992Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: 0 Warning: trajectory cut off by epoch at 20 steps.\n",
      "Epoch: 1 Warning: trajectory cut off by epoch at 9 steps.\n",
      "Epoch: 2 Warning: trajectory cut off by epoch at 38 steps.\n",
      "Epoch: 3 Warning: trajectory cut off by epoch at 7 steps.\n",
      "Epoch: 4 Warning: trajectory cut off by epoch at 32 steps.\n",
      "Epoch: 5 Warning: trajectory cut off by epoch at 5 steps.\n",
      "Epoch: 6 Warning: trajectory cut off by epoch at 56 steps.\n",
      "Epoch: 7 Warning: trajectory cut off by epoch at 119 steps.\n",
      "Epoch: 8 Warning: trajectory cut off by epoch at 12 steps.\n",
      "Epoch: 9 Warning: trajectory cut off by epoch at 141 steps.\n",
      "Epoch: 10 Warning: trajectory cut off by epoch at 53 steps.\n",
      "Epoch: 11 Warning: trajectory cut off by epoch at 19 steps.\n",
      "Epoch: 12 Warning: trajectory cut off by epoch at 132 steps.\n",
      "Epoch: 13 Warning: trajectory cut off by epoch at 69 steps.\n",
      "Epoch: 14 Warning: trajectory cut off by epoch at 198 steps.\n",
      "Epoch: 15 Warning: trajectory cut off by epoch at 48 steps.\n",
      "Epoch: 19 Warning: trajectory cut off by epoch at 107 steps.\n",
      "Epoch: 21 Warning: trajectory cut off by epoch at 89 steps.\n",
      "Epoch: 36 Warning: trajectory cut off by epoch at 113 steps.\n"
     ]
    }
   ],
   "source": [
    "# Initialize the environment and get the dimensionality of the\n",
    "# observation space and the number of possible actions\n",
    "env = gym.make(\"CartPole-v0\")\n",
    "mpi_fork(num_cpu)\n",
    "\n",
    "ppo(lambda : env, actor_critic=MLPActorCritic,\n",
    "        ac_kwargs=dict(hidden_sizes=hidden_sizes), gamma=gamma, \n",
    "        seed=702, steps_per_epoch=steps_per_epoch, epochs=epochs)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "decimal-poland",
   "metadata": {},
   "source": [
    "### visualization"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "norwegian-trouble",
   "metadata": {},
   "source": [
    "<img src='https://i.imgur.com/tKhTEaF.gif' width='600'> "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "experimental-reconstruction",
   "metadata": {},
   "source": [
    "---"
   ]
  }
 ],
 "metadata": {
  "hide_input": false,
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.9"
  },
  "toc": {
   "base_numbering": 1,
   "nav_menu": {},
   "number_sections": true,
   "sideBar": true,
   "skip_h1_title": false,
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": false,
   "toc_position": {
    "height": "calc(100% - 180px)",
    "left": "10px",
    "top": "150px",
    "width": "344px"
   },
   "toc_section_display": true,
   "toc_window_display": true
  },
  "varInspector": {
   "cols": {
    "lenName": 16,
    "lenType": 16,
    "lenVar": 40
   },
   "kernels_config": {
    "python": {
     "delete_cmd_postfix": "",
     "delete_cmd_prefix": "del ",
     "library": "var_list.py",
     "varRefreshCmd": "print(var_dic_list())"
    },
    "r": {
     "delete_cmd_postfix": ") ",
     "delete_cmd_prefix": "rm(",
     "library": "var_list.r",
     "varRefreshCmd": "cat(var_dic_list()) "
    }
   },
   "position": {
    "height": "508px",
    "left": "1347px",
    "right": "20px",
    "top": "120px",
    "width": "354px"
   },
   "types_to_exclude": [
    "module",
    "function",
    "builtin_function_or_method",
    "instance",
    "_Feature"
   ],
   "window_display": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}