Deep-Deterministic-Policy-Gradient-with-PER.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import random\n",
    "import numpy as np\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "import torch.nn.functional as F\n",
    "import torchvision.transforms as T\n",
    "from torch.nn.init import kaiming_uniform_"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import gym\n",
    "from tqdm import tqdm\n",
    "import random as rand\n",
    "from itertools import count"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class OU_noise():\n",
    "    def __init__(self,mu,sigma=0.3,theta=0.15,dt=1e-5,x0=None):\n",
    "        self.theta = theta\n",
    "        self.mu    = mu\n",
    "        self.sigma = sigma\n",
    "        self.dt    = dt\n",
    "        self.x0    = x0\n",
    "        self.reset()\n",
    "        \n",
    "    def __call__(self):\n",
    "        x = self.x_prev + self.theta * (self.mu - self.x_prev)* \\\n",
    "                self.dt +self.sigma * np.sqrt(self.dt)*np.random.normal(size = self.mu.shape)\n",
    "        self.x_prev = x\n",
    "        return x\n",
    "    \n",
    "    def reset(self):\n",
    "        self.x_prev = self.x0 if self.x0 is not None else np.zeros_like(self.mu)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class ReplayMemory():\n",
    "    def __init__(self,capacity):   \n",
    "        self.capacity = capacity\n",
    "        self.memory = []\n",
    "        self.push_count = 0\n",
    "        \n",
    "    def push(self, experience):\n",
    "        if len(self.memory) < self.capacity:\n",
    "            self.memory.append(experience)\n",
    "        else:\n",
    "            self.memory[self.push_count%self.capacity] = experience\n",
    "        self.push_count+=1\n",
    "    \n",
    "    def sample(self, batch_size):\n",
    "        return rand.sample(self.memory,batch_size)\n",
    "    \n",
    "    def can_provide_sample(self, batch_size):\n",
    "        return len(self.memory)>=batch_size\n",
    "    \n",
    "    def update_td_error(self, sampled_experiences):\n",
    "        for sampled_idx,sampled_exp in enumerate(sampled_experiences):\n",
    "            for mem_idx, mem_exp in enumerate(self.memory):\n",
    "                if mem_exp.timestep == sampled_exp.timestep:\n",
    "                    self.memory[mem_idx] = sampled_exp #update memory\n",
    "                    break\n",
    "        \n",
    "    def get_memory_values(self):\n",
    "        return self.memory    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def extract_tensors(experiences):\n",
    "    #print(\".....................................................\")\n",
    "    #print(experiences)\n",
    "    batch = Xp(*zip(*experiences))\n",
    "    state = np.stack(batch.state) #stack\n",
    "    action = np.stack(batch.action)\n",
    "    next_state = np.stack(batch.next_state)\n",
    "    reward = np.stack(batch.reward)\n",
    "    done = np.stack(batch.done)\n",
    "    abs_td_error = np.stack(batch.abs_td_error)\n",
    "    timestep = np.stack(batch.timestep)\n",
    "    return state,action,next_state,reward,done,abs_td_error,timestep"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def rebuild_experiences(state, action, next_state, reward, done, abs_error, timestep):\n",
    "    exp_list = []\n",
    "    for idx_ in range(len(state)):\n",
    "        exp_list.append(\\\n",
    "                        Xp(state[idx_], action[idx_], next_state[idx_], reward[idx_],\\\n",
    "                           done[idx_], abs_error[idx_], timestep[idx_]))\n",
    "    return exp_list"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from collections import namedtuple\n",
    "Xp = namedtuple('Experience',\n",
    "                        ('state', 'action', 'next_state', 'reward', 'done', 'abs_td_error','timestep'))\n",
    "Xp_points = Xp(5,6,7,8,9,10,11)\n",
    "Xp_points"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class linearApproximator_FCDCAP(nn.Module):\n",
    "    def __init__(self,state_shape,outputs,hidden_dims=(32,32)):\n",
    "        super(linearApproximator_FCDCAP, self).__init__()\n",
    "        self.input_size = state_shape\n",
    "        self.out = outputs\n",
    "        self.device = torch.device(\"cuda\" if torch.cuda.is_available()\\\n",
    "                                   else \"cpu\")\n",
    "        \n",
    "        self.fc1  = nn.Linear(self.input_size,hidden_dims[0])\n",
    "        self.hidden_layers = nn.ModuleList()\n",
    "        for i in range(len(hidden_dims)-1):\n",
    "            hidden_layer = nn.Linear(\\\n",
    "                                hidden_dims[i], hidden_dims[i+1])\n",
    "            self.hidden_layers.append(hidden_layer)\n",
    "        \n",
    "        self.output_layer  = nn.Linear(hidden_dims[-1],self.out)\n",
    "        self.tan_h = torch.nn.Tanh()\n",
    "        self.to(self.device)\n",
    "        \n",
    "    def forward(self, state_shape):\n",
    "        if not isinstance(state_shape, torch.Tensor):\n",
    "            state_shape = torch.tensor(state_shape, dtype=torch.float32)\n",
    "        state_shape = state_shape.to(self.device)\n",
    "            \n",
    "        x = self.fc1(state_shape)\n",
    "        x = F.relu(x)\n",
    "        \n",
    "        for hidden_layer in self.hidden_layers:\n",
    "            x = F.relu(hidden_layer(x))\n",
    "        \n",
    "        unscaled_output = self.output_layer(x)#logits, preferences of actions\n",
    "        scaled_output = self.tan_h(unscaled_output)*2\n",
    "        return scaled_output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class linearApproximator_FCQV(nn.Module):#Q value of state action pair\n",
    "    def __init__(self,state_shape,action_outputs_size,hidden_dims=(32,32)):\n",
    "        super(linearApproximator_FCQV, self).__init__()\n",
    "        self.input_size = state_shape\n",
    "        self.action_outputs_size = action_outputs_size\n",
    "        self.device = torch.device(\"cuda\" if torch.cuda.is_available()\\\n",
    "                                   else \"cpu\")\n",
    "        \n",
    "        self.fc1  = nn.Linear(self.input_size,hidden_dims[0])\n",
    "        self.hidden_layers = nn.ModuleList()\n",
    "        for i in range(len(hidden_dims)-1):\n",
    "            hidden_input_layer = hidden_dims[i]\n",
    "            if i == 0:\n",
    "                hidden_input_layer += self.action_outputs_size #increased to account for size/number of actions\n",
    "            hidden_layer = nn.Linear(\\\n",
    "                                hidden_input_layer, hidden_dims[i+1])\n",
    "            self.hidden_layers.append(hidden_layer)\n",
    "        \n",
    "        self.output_layer  = nn.Linear(hidden_dims[-1],1)\n",
    "        self.tan_h = torch.nn.Tanh()\n",
    "        self.to(self.device)\n",
    "        \n",
    "    def forward(self, state_shape, action_shape):\n",
    "        if not isinstance(state_shape, torch.Tensor):\n",
    "            state_shape = torch.tensor(state_shape, dtype=torch.float32).to(self.device)\n",
    "        if not isinstance(action_shape, torch.Tensor):\n",
    "            action_shape = torch.tensor(action_shape, dtype=torch.float32).to(self.device)\n",
    "                    \n",
    "        x = self.fc1(state_shape)\n",
    "        x = F.relu(x)\n",
    "        \n",
    "        for idx, hidden_layer in enumerate(self.hidden_layers):\n",
    "            if idx == 0:\n",
    "                x = torch.cat((x, action_shape), dim=1)\n",
    "            x = F.relu(hidden_layer(x))\n",
    "        \n",
    "        q_value = self.output_layer(x)#logits, preferences of actions\n",
    "        return q_value"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def update_networks(online_policy_network, online_value_network,\\\n",
    "                    offline_policy_network, offline_value_network, tau):\n",
    "    #polyak averaging: pg:319\n",
    "    for target_weights, online_weights in zip(offline_policy_network.parameters(), online_policy_network.parameters()):\n",
    "        target_weight_update = (1.0 - tau)*target_weights.data\n",
    "        online_weight_update = tau*online_weights.data\n",
    "        sum_up = target_weight_update + online_weight_update\n",
    "        target_weights.data.copy_(sum_up)\n",
    "        \n",
    "    for target_weights, online_weights in zip(offline_value_network.parameters(), online_value_network.parameters()):\n",
    "        target_weight_update = (1.0 - tau)*target_weights.data\n",
    "        online_weight_update = tau*online_weights.data\n",
    "        sum_up = target_weight_update + online_weight_update\n",
    "        target_weights.data.copy_(sum_up)\n",
    "    return offline_policy_network, offline_value_network"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def update_online_model(experience_samples,\\\n",
    "                        online_policy_network, online_q_network,\\\n",
    "                        offline_policy_network, offline_q_network,\\\n",
    "                        online_policy_optimizer,online_value_optimizer,\\\n",
    "                        gamma, weighted_importance, indices):\n",
    "    \n",
    "    states, actions, next_states, rewards, done, td_errors, timesteps = extract_tensors(experience_samples)\n",
    "    \n",
    "    #rearrange first\n",
    "    indices = np.array(indices)#have to be same size as batch size(already taken care of)\n",
    "    indices = np.squeeze(indices)\n",
    "    \n",
    "    states = states[indices,:]\n",
    "    states = np.squeeze(states)\n",
    "    \n",
    "    next_states = next_states[indices,:]\n",
    "    next_states = np.squeeze(next_states)\n",
    "    \n",
    "    actions = actions[indices]\n",
    "    rewards = rewards[indices]\n",
    "    done = done[indices]\n",
    "    td_errors = td_errors[indices]\n",
    "    timesteps = timesteps[indices]    \n",
    "    \n",
    "    \n",
    "    states = torch.tensor(states).float()\n",
    "    actions = torch.tensor(actions)\n",
    "    actions = actions.type(torch.LongTensor)\n",
    "    next_states=torch.tensor(next_states).float()\n",
    "    rewards = torch.tensor(rewards).float()\n",
    "    rewards = rewards.unsqueeze(1)\n",
    "    done = torch.tensor(done).float()\n",
    "    done = done.unsqueeze(1)\n",
    "    weighted_importance = torch.tensor(weighted_importance).float()\n",
    "    \n",
    "    \n",
    "    q_offline_max_actions = offline_policy_network(next_states)#get \n",
    "    q_target_next_states_action = offline_q_network(next_states,\\\n",
    "                                                    q_offline_max_actions.detach())\n",
    "    q_target_next_states_action *=(1 - done) #set terminal q vals to zero]\n",
    "    q_target = rewards + (gamma*q_target_next_states_action.detach())\n",
    "    #print(\"qt: \", q_target.shape)\n",
    "    #q_target = q_target.detach()\n",
    "    \n",
    "    q_online_state_action_val = online_q_network(states, actions)\n",
    "    #print(\"q_sa: \", q_online_state_action_val.shape)\n",
    "    q_online_value_loss = torch.nn.SmoothL1Loss()\n",
    "    q_online_state_action_val*=weighted_importance\n",
    "    #print(\"q_sa prime: \", q_online_state_action_val.shape)\n",
    "    q_target*=weighted_importance\n",
    "    q_online_value_loss = q_online_value_loss(q_online_state_action_val,\\\n",
    "                                             q_target)\n",
    "    \n",
    "    absolute_error = abs(q_online_state_action_val.detach().cpu().numpy()\\\n",
    "                         - q_target.detach().cpu().numpy())\n",
    "    absolute_error = np.squeeze(absolute_error)#used to update errors in memory\n",
    "    \n",
    "    online_value_optimizer.zero_grad()\n",
    "    q_online_value_loss.backward()\n",
    "    online_value_optimizer.step()\n",
    "    \n",
    "    #update policy_network\n",
    "    q_online_max =online_policy_network(states)\n",
    "    q_target = online_q_network(states,q_online_max)\n",
    "    q_target*=weighted_importance\n",
    "    online_policy_loss = -q_target.mean()\n",
    "    \n",
    "    online_policy_optimizer.zero_grad()\n",
    "    online_policy_loss.backward()\n",
    "    online_policy_optimizer.step()\n",
    "    \n",
    "    states, actions, next_states, rewards, done, td_errors, timesteps = extract_tensors(experience_samples)\n",
    "    #print(experience_samples)\n",
    "    experiences_rebuilded = rebuild_experiences(states, actions, next_states, rewards, done, absolute_error, timesteps)\n",
    "    #print(experiences_rebuilded)\n",
    "    return experiences_rebuilded"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def query_error(offline_policy_network, offline_q_network,\\\n",
    "                online_q_network, state, next_state, action, reward, gamma):\n",
    "    state = torch.tensor(state).float()\n",
    "    state = state.unsqueeze(0)\n",
    "    next_state = torch.tensor(next_state).float()\n",
    "    next_state = next_state.unsqueeze(0)\n",
    "    \n",
    "    q_offline_max_actions = offline_policy_network(next_state)\n",
    "    \n",
    "    q_target_next_states_action = offline_q_network(next_state,\\\n",
    "                                                    q_offline_max_actions)\n",
    "    q_target = reward + (gamma*q_target_next_states_action.detach())\n",
    "    action = np.expand_dims(action, axis=0)\n",
    "    #print(\"action after: \", action.shape)\n",
    "    q_online_state_action_val = online_q_network(state, action)\n",
    "    \n",
    "    abs_error = abs(q_target - q_online_state_action_val)\n",
    "    abs_error = np.squeeze(abs_error)\n",
    "    #print(\"abs: \",abs_error )\n",
    "    return abs_error.detach().cpu().numpy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def prioritize_samples(experience_samples, alpha, beta):\n",
    "    state,action,next_state,reward,done,abs_td_error,timesteps = extract_tensors(experience_samples)\n",
    "    #rank based\n",
    "    #('state', 'action', 'next_state', 'reward', 'done', 'abs_td_error','timestep')\n",
    "    abs_td_error  = np.expand_dims(abs_td_error, axis=1)\n",
    "    abs_td_error = torch.tensor(abs_td_error)\n",
    "    abs_td_error, indices_ = abs_td_error.sort(0, descending=True)#big to small\n",
    "    indices = np.arange(1, len(abs_td_error)+1)\n",
    "    priorities = 1.0/indices\n",
    "    priorities = priorities**alpha#scale by alpha\n",
    "    priorities = np.expand_dims(priorities, axis=1)\n",
    "    probabilities = priorities/np.sum(priorities, axis=0)#sums up to 1(or 0.999999)\n",
    "    assert np.isclose(probabilities.sum(), 1.0)#ensures probs add up to 1\n",
    "    \n",
    "    number_of_samples  = len(probabilities)\n",
    "    weight_importance_ = number_of_samples*probabilities\n",
    "    weight_importance_ = weight_importance_**-beta\n",
    "    weight_importance_max = np.max(weight_importance_)\n",
    "    weight_importance_scaled = weight_importance_/weight_importance_max\n",
    "    return weight_importance_scaled, indices_ #return weight important samples, return indices for re_arranging sampled experiences"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def freeze_model(model):\n",
    "    for param in model.parameters():\n",
    "        param.requires_grad = False\n",
    "    return model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def select_action(state, online_policy_network, noise):\n",
    "    device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "    state = torch.tensor(state).float().to(device)\n",
    "    with torch.no_grad():\n",
    "        max_action_value = online_policy_network(state).cpu().detach()\n",
    "        max_action_value = max_action_value.data.numpy().squeeze()\n",
    "        max_action_value = max_action_value + noise()\n",
    "    return max_action_value"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import time"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def DDPG_PER(env,\n",
    "         gamma=0.99,\n",
    "         alpha_pr=0.6,\n",
    "         beta_pr=0.3,\n",
    "         memory_size = 5000,\n",
    "         tau = 0.1,\n",
    "         offline_update = 100,\n",
    "         min_sample_size=320,\n",
    "         batch_size = 30,\n",
    "         n_ep=200,\n",
    "         noise=1,\n",
    "         max_steps = 100000\n",
    "         ):\n",
    "    \n",
    "    \n",
    "    observation_space = len(env.reset())\n",
    "    \n",
    "    n_actions = 1\n",
    "    online_policy_network = linearApproximator_FCDCAP(observation_space,n_actions,\\\n",
    "                                     hidden_dims=(128,64))\n",
    "    online_q_network = linearApproximator_FCQV(observation_space,\\\n",
    "                                     n_actions,hidden_dims=(128,64))\n",
    "    \n",
    "    offline_policy_network = linearApproximator_FCDCAP(observation_space,n_actions,\\\n",
    "                                     hidden_dims=(128,64))\n",
    "    offline_q_network = linearApproximator_FCQV(observation_space,\\\n",
    "                                     n_actions,hidden_dims=(128,64))\n",
    "    \n",
    "    noise_ = OU_noise(np.zeros((n_actions)))\n",
    "    offline_policy_network.eval()\n",
    "    offline_policy_network = freeze_model(offline_policy_network)\n",
    "    offline_q_network.eval()\n",
    "    offline_q_network = freeze_model(offline_q_network)\n",
    "    \n",
    "    \n",
    "    online_policy_optimizer    = torch.optim.Adam(online_policy_network.parameters(),lr=0.0008)\n",
    "    online_q_network_optimizer = torch.optim.Adam(online_q_network.parameters(),lr=0.0008)\n",
    "    \n",
    "    memory = ReplayMemory(memory_size)\n",
    "    \n",
    "    t_step = 0 #important\n",
    "    reward_per_ep = []\n",
    "    \n",
    "    \n",
    "    for e in tqdm(range(n_ep)):\n",
    "        state = env.reset()\n",
    "        reward_accumulated = 0\n",
    "        \n",
    "        while True:\n",
    "            env.render()\n",
    "            action = select_action(state, online_policy_network, noise_)\n",
    "            \n",
    "            next_state, reward, done, info = env.step(action)\n",
    "            td_error = query_error(offline_policy_network, offline_q_network,\\\n",
    "                            online_q_network, state, next_state, action, reward, gamma)\n",
    "            \n",
    "            reward_accumulated+=reward\n",
    "            is_truncated = 'TimeLimit.truncated' in info and\\\n",
    "                                info['TimeLimit.truncated']\n",
    "            is_failure = done and not is_truncated\n",
    "           \n",
    "            memory.push(Xp(state, action, next_state, reward, is_failure, td_error, t_step))\n",
    "            state = next_state\n",
    "            t_step+=1\n",
    "            if memory.can_provide_sample(min_sample_size):\n",
    "                experience_samples = memory.sample(batch_size)\n",
    "                weighted_importance, indices = prioritize_samples(experience_samples, alpha_pr, beta_pr)\n",
    "                rebuilded_exp = update_online_model(experience_samples,\\\n",
    "                                    online_policy_network, online_q_network,\\\n",
    "                                    offline_policy_network, offline_q_network,\\\n",
    "                                    online_policy_optimizer, online_q_network_optimizer,\\\n",
    "                                    gamma, weighted_importance, indices)\n",
    "                memory.update_td_error(rebuilded_exp)\n",
    "                \n",
    "            if t_step%offline_update == 0:\n",
    "                offline_policy_network, offline_q_network = update_networks(online_policy_network, online_q_network,\\\n",
    "                                                                offline_policy_network, offline_q_network, tau)\n",
    "            if done == True:\n",
    "                reward_per_ep.append(reward_accumulated)\n",
    "                #print(reward_accumulated)\n",
    "                break\n",
    "            if t_step > max_steps:\n",
    "                env.close()\n",
    "                return reward_per_ep\n",
    "    env.close()           \n",
    "    return reward_per_ep"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#env = gym.make(\"MountainCarContinuous-v0\")\n",
    "env = gym.make(\"Pendulum-v0\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "env.action_space"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "env.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rewards = DDPG_PER(env)  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from scipy.ndimage.filters import uniform_filter1d\n",
    "arr = uniform_filter1d(rewards, size=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from matplotlib import pyplot as plt\n",
    "x = []\n",
    "y = []\n",
    "for i,e in enumerate(arr):\n",
    "    y.append(i)\n",
    "    x.append(e)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.plot(y, x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "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.7.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}