IQN-DQN.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import numpy as np\n",
    "import torch.optim as optim\n",
    "from torch.nn.utils import clip_grad_norm_\n",
    "import random\n",
    "import math\n",
    "from torch.utils.tensorboard import SummaryWriter\n",
    "from collections import deque, namedtuple\n",
    "import time\n",
    "import gym\n",
    "\n",
    "\n",
    "def weight_init(layers):\n",
    "    for layer in layers:\n",
    "        torch.nn.init.kaiming_normal_(layer.weight, nonlinearity='relu')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "class IQN(nn.Module):\n",
    "    def __init__(self, state_size, action_size,layer_size, n_step, seed, layer_type=\"ff\"):\n",
    "        super(IQN, self).__init__()\n",
    "        self.seed = torch.manual_seed(seed)\n",
    "        self.input_shape = state_size\n",
    "        self.action_size = action_size\n",
    "        self.K = 32\n",
    "        self.N = 8\n",
    "        self.n_cos = 64\n",
    "        self.layer_size = layer_size\n",
    "        self.pis = torch.FloatTensor([np.pi*i for i in range(self.n_cos)]).view(1,1,self.n_cos).to(device) # Starting from 0 as in the paper \n",
    "\n",
    "        self.head = nn.Linear(self.input_shape[0], layer_size) # cound be a cnn \n",
    "        self.cos_embedding = nn.Linear(self.n_cos, layer_size)\n",
    "        self.ff_1 = nn.Linear(layer_size, layer_size)\n",
    "        self.ff_2 = nn.Linear(layer_size, action_size)\n",
    "        #weight_init([self.head_1, self.ff_1])\n",
    "\n",
    "\n",
    "        \n",
    "    def calc_cos(self, batch_size, n_tau=8):\n",
    "        \"\"\"\n",
    "        Calculating the cosinus values depending on the number of tau samples\n",
    "        \"\"\"\n",
    "        taus = torch.rand(batch_size, n_tau).to(device).unsqueeze(-1) #(batch_size, n_tau, 1)\n",
    "        cos = torch.cos(taus*self.pis)\n",
    "\n",
    "        assert cos.shape == (batch_size,n_tau,self.n_cos), \"cos shape is incorrect\"\n",
    "        return cos, taus\n",
    "    \n",
    "    def forward(self, input, num_tau=8):\n",
    "        \"\"\"\n",
    "        Quantile Calculation depending on the number of tau\n",
    "        \n",
    "        Return:\n",
    "        quantiles [ shape of (batch_size, num_tau, action_size)]\n",
    "        taus [shape of ((batch_size, num_tau, 1))]\n",
    "        \n",
    "        \"\"\"\n",
    "        batch_size = input.shape[0]\n",
    "        \n",
    "        x = torch.relu(self.head(input))\n",
    "        cos, taus = self.calc_cos(batch_size, num_tau) # cos shape (batch, num_tau, layer_size)\n",
    "        cos = cos.view(batch_size*num_tau, self.n_cos)\n",
    "        cos_x = torch.relu(self.cos_embedding(cos)).view(batch_size, num_tau, self.layer_size) # (batch, n_tau, layer)\n",
    "        \n",
    "        # x has shape (batch, layer_size) for multiplication –> reshape to (batch, 1, layer)\n",
    "        x = (x.unsqueeze(1)*cos_x).view(batch_size*num_tau, self.layer_size)\n",
    "        \n",
    "        x = torch.relu(self.ff_1(x))\n",
    "        out = self.ff_2(x)\n",
    "        \n",
    "        return out.view(batch_size, num_tau, self.action_size), taus\n",
    "    \n",
    "    def get_action(self, inputs):\n",
    "        quantiles, _ = self.forward(inputs, self.K)\n",
    "        actions = quantiles.mean(dim=1)\n",
    "        return actions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "class ReplayBuffer:\n",
    "    \"\"\"Fixed-size buffer to store experience tuples.\"\"\"\n",
    "\n",
    "    def __init__(self, buffer_size, batch_size, device, seed, gamma, n_step=1):\n",
    "        \"\"\"Initialize a ReplayBuffer object.\n",
    "        Params\n",
    "        ======\n",
    "            buffer_size (int): maximum size of buffer\n",
    "            batch_size (int): size of each training batch\n",
    "            seed (int): random seed\n",
    "        \"\"\"\n",
    "        self.device = device\n",
    "        self.memory = deque(maxlen=buffer_size)  \n",
    "        self.batch_size = batch_size\n",
    "        self.experience = namedtuple(\"Experience\", field_names=[\"state\", \"action\", \"reward\", \"next_state\", \"done\"])\n",
    "        self.seed = random.seed(seed)\n",
    "        self.gamma = gamma\n",
    "        self.n_step = n_step\n",
    "        self.n_step_buffer = deque(maxlen=self.n_step)\n",
    "    \n",
    "    def add(self, state, action, reward, next_state, done):\n",
    "        \"\"\"Add a new experience to memory.\"\"\"\n",
    "        #print(\"before:\", state,action,reward,next_state, done)\n",
    "        self.n_step_buffer.append((state, action, reward, next_state, done))\n",
    "        if len(self.n_step_buffer) == self.n_step:\n",
    "            state, action, reward, next_state, done = self.calc_multistep_return()\n",
    "            #print(\"after:\",state,action,reward,next_state, done)\n",
    "            e = self.experience(state, action, reward, next_state, done)\n",
    "            self.memory.append(e)\n",
    "    \n",
    "    def calc_multistep_return(self):\n",
    "        Return = 0\n",
    "        for idx in range(self.n_step):\n",
    "            Return += self.gamma**idx * self.n_step_buffer[idx][2]\n",
    "        \n",
    "        return self.n_step_buffer[0][0], self.n_step_buffer[0][1], Return, self.n_step_buffer[-1][3], self.n_step_buffer[-1][4]\n",
    "        \n",
    "    \n",
    "    \n",
    "    def sample(self):\n",
    "        \"\"\"Randomly sample a batch of experiences from memory.\"\"\"\n",
    "        experiences = random.sample(self.memory, k=self.batch_size)\n",
    "\n",
    "        states = torch.from_numpy(np.stack([e.state for e in experiences if e is not None])).float().to(self.device)\n",
    "        actions = torch.from_numpy(np.vstack([e.action for e in experiences if e is not None])).long().to(self.device)\n",
    "        rewards = torch.from_numpy(np.vstack([e.reward for e in experiences if e is not None])).float().to(self.device)\n",
    "        next_states = torch.from_numpy(np.stack([e.next_state for e in experiences if e is not None])).float().to(self.device)\n",
    "        dones = torch.from_numpy(np.vstack([e.done for e in experiences if e is not None]).astype(np.uint8)).float().to(self.device)\n",
    "  \n",
    "        return (states, actions, rewards, next_states, dones)\n",
    "\n",
    "    def __len__(self):\n",
    "        \"\"\"Return the current size of internal memory.\"\"\"\n",
    "        return len(self.memory)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "class DQN_Agent():\n",
    "    \"\"\"Interacts with and learns from the environment.\"\"\"\n",
    "\n",
    "    def __init__(self,\n",
    "                 state_size,\n",
    "                 action_size,\n",
    "                 layer_size,\n",
    "                 n_step,\n",
    "                 BATCH_SIZE,\n",
    "                 BUFFER_SIZE,\n",
    "                 LR,\n",
    "                 TAU,\n",
    "                 GAMMA,\n",
    "                 UPDATE_EVERY,\n",
    "                 device,\n",
    "                 seed):\n",
    "        \"\"\"Initialize an Agent object.\n",
    "        \n",
    "        Params\n",
    "        ======\n",
    "            state_size (int): dimension of each state\n",
    "            action_size (int): dimension of each action\n",
    "            layer_size (int): size of the hidden layer\n",
    "            BATCH_SIZE (int): size of the training batch\n",
    "            BUFFER_SIZE (int): size of the replay memory\n",
    "            LR (float): learning rate\n",
    "            TAU (float): tau for soft updating the network weights\n",
    "            GAMMA (float): discount factor\n",
    "            UPDATE_EVERY (int): update frequency\n",
    "            device (str): device that is used for the compute\n",
    "            seed (int): random seed\n",
    "        \"\"\"\n",
    "        self.state_size = state_size\n",
    "        self.action_size = action_size\n",
    "        self.seed = random.seed(seed)\n",
    "        self.device = device\n",
    "        self.TAU = TAU\n",
    "        self.GAMMA = GAMMA\n",
    "        self.UPDATE_EVERY = UPDATE_EVERY\n",
    "        self.BATCH_SIZE = BATCH_SIZE\n",
    "        self.Q_updates = 0\n",
    "        self.n_step = n_step\n",
    "\n",
    "        self.action_step = 4\n",
    "        self.last_action = None\n",
    "\n",
    "        # IQN-Network\n",
    "        self.qnetwork_local = IQN(state_size, action_size,layer_size, n_step, seed).to(device)\n",
    "        self.qnetwork_target = IQN(state_size, action_size,layer_size, n_step, seed).to(device)\n",
    "\n",
    "        self.optimizer = optim.Adam(self.qnetwork_local.parameters(), lr=LR)\n",
    "        print(self.qnetwork_local)\n",
    "        \n",
    "        # Replay memory\n",
    "        self.memory = ReplayBuffer(BUFFER_SIZE, BATCH_SIZE, self.device, seed, self.GAMMA, n_step)\n",
    "        \n",
    "        # Initialize time step (for updating every UPDATE_EVERY steps)\n",
    "        self.t_step = 0\n",
    "    \n",
    "    def step(self, state, action, reward, next_state, done, writer):\n",
    "        # Save experience in replay memory\n",
    "        self.memory.add(state, action, reward, next_state, done)\n",
    "        \n",
    "        # Learn every UPDATE_EVERY time steps.\n",
    "        self.t_step = (self.t_step + 1) % self.UPDATE_EVERY\n",
    "        if self.t_step == 0:\n",
    "            # If enough samples are available in memory, get random subset and learn\n",
    "            if len(self.memory) > self.BATCH_SIZE:\n",
    "                experiences = self.memory.sample()\n",
    "                loss = self.learn(experiences)\n",
    "                self.Q_updates += 1\n",
    "                writer.add_scalar(\"Q_loss\", loss, self.Q_updates)\n",
    "\n",
    "    def act(self, state, eps=0.):\n",
    "        \"\"\"Returns actions for given state as per current policy. Acting only every 4 frames!\n",
    "        \n",
    "        Params\n",
    "        ======\n",
    "            frame: to adjust epsilon\n",
    "            state (array_like): current state\n",
    "            \n",
    "        \"\"\"\n",
    "\n",
    "        if self.action_step == 4:\n",
    "            state = np.array(state)\n",
    "\n",
    "            state = torch.from_numpy(state).float().unsqueeze(0).to(self.device)\n",
    "            self.qnetwork_local.eval()\n",
    "            with torch.no_grad():\n",
    "                action_values = self.qnetwork_local.get_action(state)\n",
    "            self.qnetwork_local.train()\n",
    "\n",
    "            # Epsilon-greedy action selection\n",
    "            if random.random() > eps: # select greedy action if random number is higher than epsilon or noisy network is used!\n",
    "                action = np.argmax(action_values.cpu().data.numpy())\n",
    "                self.last_action = action\n",
    "                return action\n",
    "            else:\n",
    "                action = random.choice(np.arange(self.action_size))\n",
    "                self.last_action = action \n",
    "                return action\n",
    "\n",
    "        else:\n",
    "            self.action_step += 1\n",
    "            return self.last_action\n",
    "\n",
    "\n",
    "    def learn(self, experiences):\n",
    "        \"\"\"Update value parameters using given batch of experience tuples.\n",
    "        Params\n",
    "        ======\n",
    "            experiences (Tuple[torch.Tensor]): tuple of (s, a, r, s', done) tuples \n",
    "            gamma (float): discount factor\n",
    "        \"\"\"\n",
    "        self.optimizer.zero_grad()\n",
    "        states, actions, rewards, next_states, dones = experiences\n",
    "        # Get max predicted Q values (for next states) from target model\n",
    "        Q_targets_next, _ = self.qnetwork_target(next_states)\n",
    "        Q_targets_next = Q_targets_next.detach().max(2)[0].unsqueeze(1) # (batch_size, 1, N)\n",
    "        \n",
    "        # Compute Q targets for current states \n",
    "        Q_targets = rewards.unsqueeze(-1) + (self.GAMMA**self.n_step * Q_targets_next * (1. - dones.unsqueeze(-1)))\n",
    "        # Get expected Q values from local model\n",
    "        Q_expected, taus = self.qnetwork_local(states)\n",
    "        Q_expected = Q_expected.gather(2, actions.unsqueeze(-1).expand(self.BATCH_SIZE, 8, 1))\n",
    "\n",
    "        # Quantile Huber loss\n",
    "        td_error = Q_targets - Q_expected\n",
    "        assert td_error.shape == (self.BATCH_SIZE, 8, 8), \"wrong td error shape\"\n",
    "        huber_l = calculate_huber_loss(td_error, 1.0)\n",
    "        quantil_l = abs(taus -(td_error.detach() < 0).float()) * huber_l / 1.0\n",
    "        \n",
    "        loss = quantil_l.sum(dim=1).mean(dim=1) # , keepdim=True if per weights get multipl\n",
    "        loss = loss.mean()\n",
    "\n",
    "\n",
    "        # Minimize the loss\n",
    "        loss.backward()\n",
    "        #clip_grad_norm_(self.qnetwork_local.parameters(),1)\n",
    "        self.optimizer.step()\n",
    "\n",
    "        # ------------------- update target network ------------------- #\n",
    "        self.soft_update(self.qnetwork_local, self.qnetwork_target)\n",
    "        return loss.detach().cpu().numpy()            \n",
    "\n",
    "    def soft_update(self, local_model, target_model):\n",
    "        \"\"\"Soft update model parameters.\n",
    "        θ_target = τ*θ_local + (1 - τ)*θ_target\n",
    "        Params\n",
    "        ======\n",
    "            local_model (PyTorch model): weights will be copied from\n",
    "            target_model (PyTorch model): weights will be copied to\n",
    "            tau (float): interpolation parameter \n",
    "        \"\"\"\n",
    "        for target_param, local_param in zip(target_model.parameters(), local_model.parameters()):\n",
    "            target_param.data.copy_(self.TAU*local_param.data + (1.0-self.TAU)*target_param.data)\n",
    "\n",
    "\n",
    "def calculate_huber_loss(td_errors, k=1.0):\n",
    "    \"\"\"\n",
    "    Calculate huber loss element-wisely depending on kappa k.\n",
    "    \"\"\"\n",
    "    loss = torch.where(td_errors.abs() <= k, 0.5 * td_errors.pow(2), k * (td_errors.abs() - 0.5 * k))\n",
    "    assert loss.shape == (td_errors.shape[0], 8, 8), \"huber loss has wrong shape\"\n",
    "    return loss\n",
    "    \n",
    "            \n",
    "def eval_runs(eps, frame):\n",
    "    \"\"\"\n",
    "    Makes an evaluation run with the current epsilon\n",
    "    \"\"\"\n",
    "    env = gym.make(\"CartPole-v0\")\n",
    "    reward_batch = []\n",
    "    for i in range(5):\n",
    "        state = env.reset()\n",
    "        rewards = 0\n",
    "        while True:\n",
    "            action = agent.act(state, eps)\n",
    "            state, reward, done, _ = env.step(action)\n",
    "            rewards += reward\n",
    "            if done:\n",
    "                break\n",
    "        reward_batch.append(rewards)\n",
    "        \n",
    "    writer.add_scalar(\"Reward\", np.mean(reward_batch), frame)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Using  cuda:0\n",
      "IQN(\n",
      "  (head): Linear(in_features=4, out_features=256, bias=True)\n",
      "  (cos_embedding): Linear(in_features=64, out_features=256, bias=True)\n",
      "  (ff_1): Linear(in_features=256, out_features=256, bias=True)\n",
      "  (ff_2): Linear(in_features=256, out_features=2, bias=True)\n",
      ")\n",
      "Episode 100\tFrame 4458\tAverage Score: 44.588\n",
      "Episode 200\tFrame 19474\tAverage Score: 150.166\n",
      "Episode 300\tFrame 33258\tAverage Score: 137.844\n",
      "Episode 400\tFrame 51807\tAverage Score: 185.499\n",
      "Episode 441\tFrame 59819 \tAverage Score: 195.10Training time: 4.75min\n"
     ]
    }
   ],
   "source": [
    "def run(frames=1000, eps_fixed=False, eps_frames=1e6, min_eps=0.01):\n",
    "    \"\"\"Deep Q-Learning.\n",
    "    \n",
    "    Params\n",
    "    ======\n",
    "        n_episodes (int): maximum number of training episodes\n",
    "        max_t (int): maximum number of timesteps per episode\n",
    "        eps_start (float): starting value of epsilon, for epsilon-greedy action selection\n",
    "        eps_end (float): minimum value of epsilon\n",
    "        eps_decay (float): multiplicative factor (per episode) for decreasing epsilon\n",
    "    \"\"\"\n",
    "    scores = []                        # list containing scores from each episode\n",
    "    scores_window = deque(maxlen=100)  # last 100 scores\n",
    "    output_history = []\n",
    "    frame = 0\n",
    "    if eps_fixed:\n",
    "        eps = 0\n",
    "    else:\n",
    "        eps = 1\n",
    "    eps_start = 1\n",
    "    i_episode = 1\n",
    "    state = env.reset()\n",
    "    score = 0                  \n",
    "    for frame in range(1, frames+1):\n",
    "\n",
    "        action = agent.act(state, eps)\n",
    "        next_state, reward, done, _ = env.step(action)\n",
    "        agent.step(state, action, reward, next_state, done, writer)\n",
    "        state = next_state\n",
    "        score += reward\n",
    "        # linear annealing to the min epsilon value until eps_frames and from there slowly decease epsilon to 0 until the end of training\n",
    "        if eps_fixed == False:\n",
    "            if frame < eps_frames:\n",
    "                eps = max(eps_start - (frame*(1/eps_frames)), min_eps)\n",
    "            else:\n",
    "                eps = max(min_eps - min_eps*((frame-eps_frames)/(frames-eps_frames)), 0.001)\n",
    "\n",
    "        # evaluation runs\n",
    "        if frame % 1000 == 0:\n",
    "            eval_runs(eps, frame)\n",
    "        \n",
    "        if done:\n",
    "            scores_window.append(score)       # save most recent score\n",
    "            scores.append(score)              # save most recent score\n",
    "            writer.add_scalar(\"Average100\", np.mean(scores_window), frame)\n",
    "            output_history.append(np.mean(scores_window))\n",
    "            print('\\rEpisode {}\\tFrame {} \\tAverage Score: {:.2f}'.format(i_episode, frame, np.mean(scores_window)), end=\"\")\n",
    "            if i_episode % 100 == 0:\n",
    "                print('\\rEpisode {}\\tFrame {}\\tAverage Score: {:.2f}'.format(i_episode,frame, np.mean(scores_window)))\n",
    "            i_episode +=1 \n",
    "            state = env.reset()\n",
    "            score = 0              \n",
    "\n",
    "    return output_history\n",
    "\n",
    "\n",
    "if __name__ == \"__main__\":\n",
    "\n",
    "    writer = SummaryWriter(\"runs/\"+\"IQN_CP_5\")\n",
    "    seed = 5\n",
    "    BUFFER_SIZE = 100000\n",
    "    BATCH_SIZE = 8\n",
    "    GAMMA = 0.99\n",
    "    TAU = 1e-2\n",
    "    LR = 1e-3\n",
    "    UPDATE_EVERY = 1\n",
    "    n_step = 1\n",
    "    device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n",
    "    print(\"Using \", device)\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "    np.random.seed(seed)\n",
    "    env = gym.make(\"CartPole-v0\")\n",
    "\n",
    "    env.seed(seed)\n",
    "    action_size     = env.action_space.n\n",
    "    state_size = env.observation_space.shape\n",
    "\n",
    "    agent = DQN_Agent(state_size=state_size,    \n",
    "                        action_size=action_size,\n",
    "                        layer_size=256,\n",
    "                        n_step=n_step,\n",
    "                        BATCH_SIZE=BATCH_SIZE, \n",
    "                        BUFFER_SIZE=BUFFER_SIZE, \n",
    "                        LR=LR, \n",
    "                        TAU=TAU, \n",
    "                        GAMMA=GAMMA, \n",
    "                        UPDATE_EVERY=UPDATE_EVERY, \n",
    "                        device=device, \n",
    "                        seed=seed)\n",
    "\n",
    "\n",
    "\n",
    "    # set epsilon frames to 0 so no epsilon exploration\n",
    "    eps_fixed = False\n",
    "\n",
    "    t0 = time.time()\n",
    "    final_average100 = run(frames = 60000, eps_fixed=eps_fixed, eps_frames=5000, min_eps=0.025)\n",
    "    t1 = time.time()\n",
    "    \n",
    "    print(\"Training time: {}min\".format(round((t1-t0)/60,2)))\n",
    "    torch.save(agent.qnetwork_local.state_dict(), \"IQN\"+\".pth\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "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.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}