REINFORCE_(usingMultipleEpisodeUpdates).ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "import gym\n",
    "import time\n",
    "import torch\n",
    "import random\n",
    "import numpy as np\n",
    "import torch.nn as nn\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[2020-03-30 19:03:34,047] Making new env: InvertedPendulum-v1\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Observation Dimension: 4 | Action Dimension: 1\n"
     ]
    }
   ],
   "source": [
    "# env = gym.make('HalfCheetah-v1')\n",
    "env = gym.make('InvertedPendulum-v1')\n",
    "print('Observation Dimension:', env.observation_space.shape[0], '| Action Dimension:', env.action_space.shape[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Config():\n",
    "\n",
    "    def __init__(self, epsilon_start, epsilon_end, decay, batch_size, batch_length, observation_dim, action_dim, N_episodes, discount, learning_rate):\n",
    "        self.epsilon_start = epsilon_start\n",
    "        self.epsilon_end = epsilon_end\n",
    "        self.decay = decay\n",
    "        self.batch_size = batch_size\n",
    "        self.batch_length = batch_length\n",
    "        self.observation_dim = observation_dim\n",
    "        self.action_dim = action_dim\n",
    "        self.device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "        self.N_episodes = N_episodes\n",
    "        self.discount = discount\n",
    "        self.learning_rate = learning_rate\n",
    "\n",
    "        self.rewards = []\n",
    "    \n",
    "    def append_reward(self, episodic_reward):\n",
    "        self.rewards.append(episodic_reward)\n",
    "        \n",
    "    def plot(self):\n",
    "        clear_output(True)\n",
    "        plt.figure(figsize=(20,5))\n",
    "        plt.subplot(131)\n",
    "        plt.title('Reward: %s' % (np.mean(self.rewards[-10:])))\n",
    "        plt.plot(self.rewards)\n",
    "#         plt.subplot(132)\n",
    "#         plt.title('loss')\n",
    "#         plt.plot(losses)\n",
    "        plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "# One hidden layer function approximator\n",
    "class Actor(nn.Module):\n",
    "    \"\"\" Policy returns continuous action\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, config, std=0.0):\n",
    "        super(Actor, self).__init__()\n",
    "        self.eps = config.epsilon_start\n",
    "        self.eps_end = config.epsilon_end\n",
    "        self.eps_decay = config.decay\n",
    "        self.device = config.device\n",
    "        \n",
    "        self.criterion = torch.nn.MSELoss()\n",
    "        self.batch_size = config.batch_size\n",
    "        \n",
    "        self.block = nn.Sequential(\n",
    "            nn.Linear(config.observation_dim, 16),\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(16, config.action_dim),\n",
    "        )\n",
    "        \n",
    "        self.log_std = nn.Parameter(torch.ones(1, config.action_dim) * std)\n",
    "\n",
    "\n",
    "    def forward(self, ip):\n",
    "        return self.block(ip)\n",
    "    \n",
    "    def act(self, observation):\n",
    "        self.eps = max(self.eps*self.eps_decay, self.eps_end)      # Decay Doubt Batch random acting doubt\n",
    "        observation = torch.FloatTensor(observation).to(self.device)\n",
    "\n",
    "        # Q: Do i need to use eplison greedy policy? \n",
    "        action_mean = self.forward(observation).unsqueeze(dim=0)#.squeeze()\n",
    "        \n",
    "        assert action_mean.shape == (1, 1)\n",
    "        \n",
    "        std   = self.log_std.exp().expand_as(action_mean)\n",
    "        dist = torch.distributions.Normal(action_mean, std)\n",
    "        actions = dist.sample()\n",
    "        log_probs = dist.log_prob(actions)\n",
    "        \n",
    "        return actions.detach().cpu().numpy(), log_probs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "# One hidden layer function approximator\n",
    "class Critic(nn.Module):\n",
    "    \"\"\" Policy returns continuous action\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, config):\n",
    "        super(Critic, self).__init__()\n",
    "        self.eps = config.epsilon_start\n",
    "        self.eps_end = config.epsilon_end\n",
    "        self.eps_decay = config.decay\n",
    "        self.device = config.device\n",
    "        \n",
    "        self.criterion = torch.nn.MSELoss()\n",
    "        self.batch_size = config.batch_size\n",
    "        \n",
    "        self.block = nn.Sequential(\n",
    "            nn.Linear(config.observation_dim, 64),\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(64, 32),\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(32, config.action_dim),\n",
    "        )\n",
    "\n",
    "    def forward(self, ip):\n",
    "        return self.block(ip)\n",
    "    \n",
    "    def act(self, observation):\n",
    "        self.eps = max(self.eps*self.eps_decay, self.eps_end)      # Decay Doubt Batch random acting doubt\n",
    "        observation = torch.FloatTensor(observation).to(self.device)\n",
    "\n",
    "        if random.random() < self.eps:\n",
    "            return env.action_space.sample()\n",
    "        else:\n",
    "            return self.forward(observation).squeeze().max(dim=1)[1].cpu().detach().numpy() # 1 returns index, 0 returns value"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "config = Config(epsilon_start=0.99, epsilon_end=0.1, decay=0.9999, batch_size=64, batch_length=100,\n",
    "                observation_dim=env.observation_space.shape[0], action_dim=env.action_space.shape[0],\n",
    "                N_episodes=1000, discount=0.99, learning_rate = 3e-2\n",
    "               )\n",
    "\n",
    "actor = Actor(config).to(config.device)\n",
    "critic = Critic(config).to(config.device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "from IPython.display import clear_output\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "class TrajectoryBuffer():\n",
    "    \"\"\" Policy returns continuous action\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, config):\n",
    "        self.log_probs = []\n",
    "        self.returns = []\n",
    "        self.discount = config.discount\n",
    "        \n",
    "    def clear(self):\n",
    "        self.log_probs = []\n",
    "        self.returns = []\n",
    "    \n",
    "    def push(self, log_probs, rewards):\n",
    "        \"Convert to discounted returns and push to buffer\"\n",
    "        Rt = rewards.copy()\n",
    "        discounts = [self.discount**i for i in range(len(Rt))]\n",
    "        discounts.reverse()        \n",
    "        Rt.reverse()\n",
    "        returns = [np.sum(np.multiply(Rt[:len(Rt)-i], discounts[i:])) for i in range(len(Rt))]\n",
    "        \n",
    "        self.log_probs.append(log_probs)\n",
    "        self.returns.append(returns)\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "def update_actor():\n",
    "    \n",
    "    optimizer.zero_grad()\n",
    "    loss = 0\n",
    "    for it, log_prob in enumerate(trajectory_buffer.log_probs):\n",
    "        return_ = trajectory_buffer.returns[it]\n",
    "        for it_, return__ in enumerate(return_):\n",
    "            log_prob_ = log_prob[it_]\n",
    "            loss -= return__*log_prob_\n",
    "\n",
    "    loss = loss/(it+1)\n",
    "    loss.backward()\n",
    "    optimizer.step()\n",
    "    \n",
    "    trajectory_buffer.clear()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "trajectory_buffer = TrajectoryBuffer(config)\n",
    "optimizer = torch.optim.Adam(actor.parameters(), lr=config.learning_rate)\n",
    "last_highest_rew = 0\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "for i in range(config.N_episodes):\n",
    "    \n",
    "    for _ in range(config.batch_length):\n",
    "        \n",
    "        episodic_rewards, log_probs, done = [], [], False\n",
    "        state = env.reset()\n",
    "\n",
    "        while not done:\n",
    "\n",
    "            action, log_prob = actor.act(state)\n",
    "\n",
    "            next_state, reward, done, _ = env.step(action)\n",
    "\n",
    "            episodic_rewards.append(reward)\n",
    "            log_probs.append(log_prob)\n",
    "            if i%5 == 0:\n",
    "                env.render()\n",
    "\n",
    "            state = next_state\n",
    "\n",
    "        trajectory_buffer.push(log_probs, episodic_rewards)\n",
    "\n",
    "        if np.sum(episodic_rewards) > last_highest_rew:\n",
    "            torch.save(actor.state_dict(), 'InvertedPendulumActor.pt')\n",
    "            last_highest_rew = np.sum(episodic_rewards)\n",
    "            print('Saved Model')\n",
    "\n",
    "        if np.sum(episodic_rewards) > 200:\n",
    "            print('Wohoo!! Agent Trained')\n",
    "            break\n",
    "\n",
    "        #Update\n",
    "        config.append_reward(np.sum(episodic_rewards))\n",
    "        config.plot()\n",
    "        \n",
    "    update_actor()   \n",
    "    print('##############UPDATED_ACTOR############')\n",
    "    time.sleep(2)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Actor(\n",
       "  (criterion): MSELoss()\n",
       "  (block): Sequential(\n",
       "    (0): Linear(in_features=4, out_features=16, bias=True)\n",
       "    (1): ReLU()\n",
       "    (2): Linear(in_features=16, out_features=1, bias=True)\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "actor_test = Actor(config).to(config.device)\n",
    "actor_test.load_state_dict(torch.load('InvertedPendulumActor.pt'))\n",
    "# actor_test.load_state_dict(actor.state_dict())\n",
    "actor_test.eval()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<h2> Testing policy</h2>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Total collected reward: 1000.0\n"
     ]
    }
   ],
   "source": [
    "import time\n",
    "\n",
    "\n",
    "state = env.reset()\n",
    "total_reward = 0\n",
    "done = False\n",
    "\n",
    "while not done:\n",
    "\n",
    "    action, log_prob = actor_test.act(state)\n",
    "\n",
    "    next_state, reward, done, _ = env.step(action)\n",
    "    total_reward += reward\n",
    "\n",
    "    env.render()\n",
    "    time.sleep(0.001)\n",
    "\n",
    "    state = next_state\n",
    "\n",
    "print('Total collected reward:', total_reward)\n",
    "\n",
    "env.close()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<h2> Test Ground</h2>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1000.0\n"
     ]
    }
   ],
   "source": [
    "print(last_highest_rew)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}