BDQN-Pytorch.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import argparse\n",
    "import math\n",
    "import random\n",
    "from copy import deepcopy\n",
    "from torch.distributions import normal\n",
    "\n",
    "import numpy as np\n",
    "import torch\n",
    "import torch.optim as optim\n",
    "from helpers import ReplayBuffer, make_atari, make_gym_env, wrap_deepmind, wrap_pytorch\n",
    "from models import Deep_feature, CnnDQN"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "NOT Using GPU: GPU not requested or not available.\n"
     ]
    }
   ],
   "source": [
    "USE_CUDA = torch.cuda.is_available()\n",
    "if USE_CUDA:\n",
    "    print(\"Using GPU: GPU requested and available.\")\n",
    "    dtype = torch.cuda.FloatTensor\n",
    "    dtypelong = torch.cuda.LongTensor\n",
    "else:\n",
    "    print(\"NOT Using GPU: GPU not requested or not available.\")\n",
    "    dtype = torch.FloatTensor\n",
    "    dtypelong = torch.LongTensor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "feature_dimension = 512\n",
    "learning_rate = 0.0025 \n",
    "replay_buffer_size = 1000000\n",
    "max_time_step = 5 * 10**6\n",
    "\n",
    "sigma = 0.001\n",
    "sigma_n = 1\n",
    "\n",
    "start_train_ts = 10*5\n",
    "batch_size = 32\n",
    "gamma = 0.99\n",
    "\n",
    "target_network_update_f = 100#10000\n",
    "\n",
    "target_W_update = 1 #10\n",
    "target_batch_size = 500 #5000\n",
    "\n",
    "log_every = 2000\n",
    "\n",
    "env_name = \"PongNoFrameskip-v4\"  # Set the desired environment\n",
    "env = make_atari(env_name)\n",
    "env = wrap_pytorch(wrap_deepmind(env, scale=True))\n",
    "num_action = env.action_space.n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "def compute_td_loss(batch_size, replay_buffer, optimizer):\n",
    "    state, action, reward, next_state, done = replay_buffer.sample(batch_size)\n",
    "    state = torch.tensor(np.float32(state)).type(dtype)\n",
    "    next_state = torch.tensor(np.float32(next_state)).type(dtype)\n",
    "    action = torch.tensor(action).type(dtypelong)\n",
    "    reward = torch.tensor(reward).type(dtype)\n",
    "    done = torch.tensor(done).type(dtype)\n",
    "\n",
    "    _, argmax_Q = torch.max(torch.mm(deep_feature(next_state), W_mean.transpose(0, 1)),dim=1,keepdim=True)\n",
    "    Q_target = torch.mm(deep_target_feature(next_state), W_target.transpose(0, 1))\n",
    "    Q_target = torch.gather(Q_target, 1, argmax_Q).squeeze() * (1 - done)\n",
    "    Q = torch.mm(deep_feature(state), W_mean.transpose(0, 1))\n",
    "    Q = torch.gather(Q, 1, action.type(dtypelong).unsqueeze(1)).squeeze()\n",
    "    target = (reward + gamma * Q_target).data\n",
    "    loss = (Q - target).pow(2).mean()\n",
    "    optimizer.zero_grad()\n",
    "    loss.backward()\n",
    "    optimizer.step()\n",
    "    return loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "def BayesReg(phiphiT, phiY, target_batch_size):\n",
    "    with torch.no_grad():\n",
    "        chunk_size = 10#1000\n",
    "        num_chunks = int(target_batch_size / chunk_size)\n",
    "\n",
    "        for _ in range(num_chunks):\n",
    "            state, action, reward, next_state, done = replay_buffer.sample(chunk_size)\n",
    "            _, argmax_Q = torch.max(torch.mm(deep_feature(next_state), W_mean.transpose(0, 1)),dim=1,keepdim=True)\n",
    "            Q_target = torch.mm(deep_target_feature(next_state), W_target.transpose(0, 1))\n",
    "            Q_target = torch.gather(Q_target, 1, argmax_Q).squeeze() * (1 - done)\n",
    "            target = (reward + gamma * Q_target).data\n",
    "            \n",
    "            feature_rep = deep_feature(state).unsqueeze(1).detach()\n",
    "            for i in range(num_action):\n",
    "                action_ = action == i # I am not sure it is a right way of doing it\n",
    "                print(action_)\n",
    "                feature_rep_of_action = torch.mm(feature_rep,action_)\n",
    "                phiphiT[i] = torch.mm(feature_rep_of_action.transpose(0, 1),feature_rep_of_action)\n",
    "                phiY[i] = torch.mm(feature_rep_of_action,target)\n",
    "\n",
    "        for i in range(num_action):\n",
    "            inv = np.linalg.inv(\n",
    "                ((phiphiT[i] / sigma_n + 1 / sigma * eye).cpu()).numpy()\n",
    "            )\n",
    "            W[i] = torch.tensor(np.dot(inv, phiY[0].cpu().data) / sigma_n).type(\n",
    "                dtype\n",
    "            )\n",
    "            Cov_W[i] = torch.tensor(sigma * inv).type(dtype)\n",
    "        return phiphiT, phiY, W_mean, Cov_W"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "def Sample_W(W_mean, Cov_W_decom):\n",
    "    dist = normal.Normal(loc=0, scale=1)\n",
    "    for i in range(num_action):\n",
    "        sam = dist.sample((feature_dimension, 1)).type(dtype)\n",
    "        W[i] = W_mean[i] + torch.mm(Cov_W_decom[i], sam)[:, 0]\n",
    "    return W"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "deep_feature = Deep_feature(env.observation_space.shape,feature_dimension, env.action_space.n)\n",
    "deep_target_feature = deepcopy(deep_feature)\n",
    "\n",
    "optimizer = optim.RMSprop(deep_feature.parameters(), lr=learning_rate)\n",
    "replay_buffer = ReplayBuffer(replay_buffer_size)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "eye = torch.eye(feature_dimension).type(dtype)\n",
    "dist = normal.Normal(loc=0, scale=0.01)\n",
    "W = dist.sample((num_action, feature_dimension)).type(dtype)\n",
    "W_target = dist.sample((num_action, feature_dimension)).type(dtype)\n",
    "W_mean = dist.sample((num_action, feature_dimension)).type(dtype)\n",
    "Cov_W = torch.eye(feature_dimension).repeat(num_action, 1, 1).type(dtype)\n",
    "Cov_W_decom = Cov_W\n",
    "Cov_W_target = Cov_W\n",
    "phiphiT = torch.zeros((num_action, feature_dimension, feature_dimension)).type(dtype)\n",
    "phiY = torch.zeros((num_action, feature_dimension)).type(dtype)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "True\n"
     ]
    },
    {
     "ename": "TypeError",
     "evalue": "conv2d(): argument 'input' (position 1) must be Tensor, not numpy.ndarray",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mTypeError\u001b[0m                                 Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-9-bf7016e6c975>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m     34\u001b[0m         \u001b[0;32mif\u001b[0m \u001b[0mc_t\u001b[0m \u001b[0;34m==\u001b[0m \u001b[0mtarget_W_update\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     35\u001b[0m             \u001b[0mc_t\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;36m0\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 36\u001b[0;31m             \u001b[0mphiphiT\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mphiY\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mW_mean\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mCov_W\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mBayesReg\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mphiphiT\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mphiY\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtarget_batch_size\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     37\u001b[0m             \u001b[0mW_target\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mW_mean\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     38\u001b[0m             \u001b[0mCov_W_target\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mCov_W\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m<ipython-input-5-0d3aab2022ba>\u001b[0m in \u001b[0;36mBayesReg\u001b[0;34m(phiphiT, phiY, target_batch_size)\u001b[0m\n\u001b[1;32m      6\u001b[0m         \u001b[0;32mfor\u001b[0m \u001b[0m_\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mrange\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mnum_chunks\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      7\u001b[0m             \u001b[0mstate\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0maction\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mreward\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mnext_state\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mdone\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mreplay_buffer\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0msample\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mchunk_size\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m----> 8\u001b[0;31m             \u001b[0m_\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0margmax_Q\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mmax\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mmm\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mdeep_feature\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mnext_state\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mW_mean\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtranspose\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;36m0\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;36m1\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0mdim\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;36m1\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0mkeepdim\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;32mTrue\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m      9\u001b[0m             \u001b[0mQ_target\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mmm\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mdeep_target_feature\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mnext_state\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mW_target\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtranspose\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;36m0\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;36m1\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     10\u001b[0m             \u001b[0mQ_target\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mgather\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mQ_target\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;36m1\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0margmax_Q\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0msqueeze\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;34m*\u001b[0m \u001b[0;34m(\u001b[0m\u001b[0;36m1\u001b[0m \u001b[0;34m-\u001b[0m \u001b[0mdone\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/usr/local/lib/python3.7/site-packages/torch/nn/modules/module.py\u001b[0m in \u001b[0;36m__call__\u001b[0;34m(self, *input, **kwargs)\u001b[0m\n\u001b[1;32m    475\u001b[0m             \u001b[0mresult\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_slow_forward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0minput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    476\u001b[0m         \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 477\u001b[0;31m             \u001b[0mresult\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mforward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0minput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    478\u001b[0m         \u001b[0;32mfor\u001b[0m \u001b[0mhook\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_forward_hooks\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mvalues\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    479\u001b[0m             \u001b[0mhook_result\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mhook\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0minput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mresult\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/Desktop/WorkSpace/RL/BDQN-Pytorch/minimal_dqn/models.py\u001b[0m in \u001b[0;36mforward\u001b[0;34m(self, x)\u001b[0m\n\u001b[1;32m     24\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     25\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0mforward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mx\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 26\u001b[0;31m         \u001b[0mx\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mfeatures\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mx\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     27\u001b[0m         \u001b[0mx\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mx\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mview\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mx\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0msize\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;36m0\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m-\u001b[0m\u001b[0;36m1\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     28\u001b[0m         \u001b[0mx\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mfc\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mx\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/usr/local/lib/python3.7/site-packages/torch/nn/modules/module.py\u001b[0m in \u001b[0;36m__call__\u001b[0;34m(self, *input, **kwargs)\u001b[0m\n\u001b[1;32m    475\u001b[0m             \u001b[0mresult\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_slow_forward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0minput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    476\u001b[0m         \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 477\u001b[0;31m             \u001b[0mresult\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mforward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0minput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    478\u001b[0m         \u001b[0;32mfor\u001b[0m \u001b[0mhook\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_forward_hooks\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mvalues\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    479\u001b[0m             \u001b[0mhook_result\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mhook\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0minput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mresult\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/usr/local/lib/python3.7/site-packages/torch/nn/modules/container.py\u001b[0m in \u001b[0;36mforward\u001b[0;34m(self, input)\u001b[0m\n\u001b[1;32m     89\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0mforward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0minput\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     90\u001b[0m         \u001b[0;32mfor\u001b[0m \u001b[0mmodule\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_modules\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mvalues\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 91\u001b[0;31m             \u001b[0minput\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mmodule\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0minput\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     92\u001b[0m         \u001b[0;32mreturn\u001b[0m \u001b[0minput\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     93\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/usr/local/lib/python3.7/site-packages/torch/nn/modules/module.py\u001b[0m in \u001b[0;36m__call__\u001b[0;34m(self, *input, **kwargs)\u001b[0m\n\u001b[1;32m    475\u001b[0m             \u001b[0mresult\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_slow_forward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0minput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    476\u001b[0m         \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 477\u001b[0;31m             \u001b[0mresult\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mforward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0minput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    478\u001b[0m         \u001b[0;32mfor\u001b[0m \u001b[0mhook\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_forward_hooks\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mvalues\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    479\u001b[0m             \u001b[0mhook_result\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mhook\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0minput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mresult\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/usr/local/lib/python3.7/site-packages/torch/nn/modules/conv.py\u001b[0m in \u001b[0;36mforward\u001b[0;34m(self, input)\u001b[0m\n\u001b[1;32m    299\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0mforward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0minput\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    300\u001b[0m         return F.conv2d(input, self.weight, self.bias, self.stride,\n\u001b[0;32m--> 301\u001b[0;31m                         self.padding, self.dilation, self.groups)\n\u001b[0m\u001b[1;32m    302\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    303\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mTypeError\u001b[0m: conv2d(): argument 'input' (position 1) must be Tensor, not numpy.ndarray"
     ]
    }
   ],
   "source": [
    "losses, all_rewards = [], []\n",
    "state = env.reset()\n",
    "c_t = 0\n",
    "episode_reward = 0\n",
    "replay_memory = ReplayBuffer(replay_buffer_size)\n",
    "\n",
    "for ts in range(1, max_time_step + 1):\n",
    "    torch_state = torch.tensor(np.float32(state)).type(dtype).unsqueeze(0)\n",
    "    action = torch.mm(W, deep_feature(torch_state).transpose(0, 1)).squeeze()\n",
    "    action = torch.argmax(action)\n",
    "    \n",
    "    next_state, reward, done, _ = env.step(int(action.cpu()))\n",
    "    replay_buffer.push(state, action, reward, next_state, done)\n",
    "\n",
    "    state = next_state\n",
    "    episode_reward += reward\n",
    "\n",
    "    if done:\n",
    "        state = env.reset()\n",
    "        all_rewards.append(episode_reward)\n",
    "        episode_reward = 0\n",
    "        W = Sample_W(W_mean, Cov_W_decom)\n",
    "\n",
    "    if ( len(replay_buffer) > start_train_ts) and (ts % target_network_update_f == 0):\n",
    "        loss = compute_td_loss(batch_size, replay_buffer, optimizer)\n",
    "        losses.append(loss.data)\n",
    "    \n",
    "    if ts % target_network_update_f == 0:\n",
    "        for t_param, param in zip(deep_target_feature.parameters(), deep_feature.parameters()):\n",
    "            new_param = param.data\n",
    "            t_param.data.copy_(new_param)\n",
    "        c_t += 1\n",
    "        print(c_t == target_W_update)\n",
    "        if c_t == target_W_update:\n",
    "            c_t = 0 \n",
    "            phiphiT, phiY, W_mean, Cov_W = BayesReg(phiphiT, phiY, target_batch_size)\n",
    "            W_target = W_mean\n",
    "            Cov_W_target = Cov_W\n",
    "\n",
    "            for ii in range(num_action):\n",
    "                Cov_W_decom[ii] = torch.tensor(\n",
    "                    np.linalg.cholesky(\n",
    "                        (((Cov_W[ii] + Cov_W[ii].transpose(0, 1))) / 2.0).cpu()\n",
    "                    )\n",
    "                ).type(dtype) # in pytorch has stable cholesky decomposinong, it is better to use it, mxnet did not have\n",
    "        \n",
    "        if len(replay_memory) < 100000:\n",
    "            target_batch_size = len(replay_memory)\n",
    "        else:\n",
    "            target_batch_size = 100000\n",
    "            \n",
    "\n",
    "    if ts % log_every == 0:\n",
    "        out_str = \"Timestep {}\".format(ts)\n",
    "        if len(all_rewards) > 0:\n",
    "            out_str += \", Reward: {}\".format(all_rewards[-1])\n",
    "        if len(losses) > 0:\n",
    "            out_str += \", TD Loss: {}\".format(losses[-1])\n",
    "        print(out_str)\n"
   ]
  },
  {
   "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.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}