model_bco-discrete.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "os.environ['CUDA_VISIBLE_DEVICES'] = '0'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "from tqdm import tqdm_notebook as tqdm\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch as T\n",
    "import torch.nn as nn"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(4,) 2\n",
      "[-0.03532913  0.02262459 -0.00976092 -0.01836328]\n",
      "(array([-0.03487663, -0.17235602, -0.01012818,  0.27122405]), 1.0, False, {})\n"
     ]
    }
   ],
   "source": [
    "import gym\n",
    "\n",
    "env = gym.make('CartPole-v0')\n",
    "print(env.observation_space.shape, env.action_space.n)\n",
    "\n",
    "print(env.reset())\n",
    "print(env.step(0))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "class BCO(nn.Module):\n",
    "    def __init__(self, env, policy='mlp'):\n",
    "        super(BCO, self).__init__()\n",
    "        \n",
    "        self.policy = policy\n",
    "        self.act_n = env.action_space.n\n",
    "        \n",
    "        if self.policy=='mlp':\n",
    "            self.obs_n = env.observation_space.shape[0]\n",
    "            self.pol = nn.Sequential(*[nn.Linear(self.obs_n, 32), nn.LeakyReLU(), \n",
    "                                       nn.Linear(32, 32), nn.LeakyReLU(), \n",
    "                                       nn.Linear(32, self.act_n)])\n",
    "            self.inv = nn.Sequential(*[nn.Linear(self.obs_n*2, 32), nn.LeakyReLU(), \n",
    "                                       nn.Linear(32, 32), nn.LeakyReLU(),  \n",
    "                                       nn.Linear(32, self.act_n)])\n",
    "        \n",
    "        elif self.policy=='cnn':\n",
    "            pass\n",
    "    \n",
    "    def pred_act(self, obs):\n",
    "        out = self.pol(obs)\n",
    "        \n",
    "        return out\n",
    "    \n",
    "    def pred_inv(self, obs1, obs2):\n",
    "        obs = T.cat([obs1, obs2], dim=1)\n",
    "        out = self.inv(obs)\n",
    "        \n",
    "        return out\n",
    "\n",
    "POLICY = 'mlp'\n",
    "model = BCO(env, policy=POLICY).cuda()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch.utils.data import Dataset, DataLoader\n",
    "\n",
    "class DS_Inv(Dataset):\n",
    "    def __init__(self, trajs):\n",
    "        self.dat = []\n",
    "        \n",
    "        for traj in trajs:\n",
    "            for dat in traj:\n",
    "                obs, act, new_obs = dat\n",
    "                \n",
    "                self.dat.append([obs, new_obs, act])\n",
    "    \n",
    "    def __len__(self):\n",
    "        return len(self.dat)\n",
    "    \n",
    "    def __getitem__(self, idx):\n",
    "        obs, new_obs, act = self.dat[idx]\n",
    "        \n",
    "        return obs, new_obs, np.asarray(act)\n",
    "\n",
    "class DS_Policy(Dataset):\n",
    "    def __init__(self, traj):\n",
    "        self.dat = []\n",
    "        \n",
    "        for dat in traj:\n",
    "            obs, act = dat\n",
    "                \n",
    "            self.dat.append([obs, act])\n",
    "    \n",
    "    def __len__(self):\n",
    "        return len(self.dat)\n",
    "    \n",
    "    def __getitem__(self, idx):\n",
    "        obs, act = self.dat[idx]\n",
    "        \n",
    "        return obs, np.asarray(act)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "10\n",
      "torch.Size([100, 4]) torch.Size([100, 4])\n"
     ]
    }
   ],
   "source": [
    "import pickle\n",
    "\n",
    "trajs_demo = pickle.load(open('Demo/demo_cart-pole.pkl', 'rb'))\n",
    "ld_demo = DataLoader(DS_Inv(trajs_demo), batch_size=100)\n",
    "\n",
    "print(len(ld_demo))\n",
    "for obs1, obs2, _ in ld_demo:\n",
    "    print(obs1.shape, obs2.shape)\n",
    "    \n",
    "    break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "loss_func = nn.CrossEntropyLoss().cuda()\n",
    "optim = T.optim.Adam(model.parameters(), lr=5e-3)\n",
    "\n",
    "EPOCHS = 15\n",
    "M = 1000\n",
    "\n",
    "EPS = 0.9\n",
    "DECAY = 0.5"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "95909ad9e3de40a6a3d6e5cb7e59b706",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(IntProgress(value=0, max=15), HTML(value='')))"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Ep 1: reward=21.21\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "589293d29ca04b1caa26ee923b5310ca",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(IntProgress(value=0, max=32), HTML(value='')))"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Ep 1: loss_inv=0.596\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "29fa5cfe59584faeb6c4e9e041b2df17",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(IntProgress(value=0, max=32), HTML(value='')))"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Ep 1: loss_policy=0.553\n",
      "Ep 2: reward=47.71\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "a354543433ff49d0a4c0cb4ee0f3d66e",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(IntProgress(value=0, max=64), HTML(value='')))"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Ep 2: loss_inv=0.023\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "fbfe933e65444977b0f67ccad0f3ddcb",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(IntProgress(value=0, max=32), HTML(value='')))"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Ep 2: loss_policy=0.545\n",
      "Ep 3: reward=50.00\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "f22f1817a2f9475db93cea327b073122",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(IntProgress(value=0, max=95), HTML(value='')))"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Ep 3: loss_inv=0.000\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "7112c58ca7dd4a8a96965604b1080f84",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(IntProgress(value=0, max=32), HTML(value='')))"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Ep 3: loss_policy=0.606\n",
      "Ep 4: reward=56.11\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "5afe976906fc443ebae9f00abf7f3288",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(IntProgress(value=0, max=126), HTML(value='')))"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Ep 4: loss_inv=0.000\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "6a511b5cd960475ea0c6a0678f62d185",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(IntProgress(value=0, max=32), HTML(value='')))"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Ep 4: loss_policy=0.460\n",
      "Ep 5: reward=44.48\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "01c980e79dc344b6a2a8a34898d42d56",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(IntProgress(value=0, max=158), HTML(value='')))"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Ep 5: loss_inv=0.000\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "2fcedea57c5246d6ba57eb9be08c9ed5",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(IntProgress(value=0, max=32), HTML(value='')))"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Ep 5: loss_policy=0.443\n",
      "Ep 6: reward=45.77\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "47424532a19c428a9bf671c9bfa00d35",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(IntProgress(value=0, max=190), HTML(value='')))"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Ep 6: loss_inv=0.000\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "e337a352ab474396929f99fdc3a1f795",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(IntProgress(value=0, max=32), HTML(value='')))"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Ep 6: loss_policy=0.428\n",
      "Ep 7: reward=54.63\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "21ba647643834a4ab18ed54dfa86fd03",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(IntProgress(value=0, max=222), HTML(value='')))"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Ep 7: loss_inv=0.000\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "fbc83b5b9d3c4a2e9b60b03ed1ea6f32",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(IntProgress(value=0, max=32), HTML(value='')))"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Ep 7: loss_policy=0.449\n",
      "Ep 8: reward=52.35\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "70ee90a09274493185b333dd44439d58",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(IntProgress(value=0, max=255), HTML(value='')))"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Ep 8: loss_inv=0.000\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "cfa6d3ff8e9043859291d24f461f1a7e",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(IntProgress(value=0, max=32), HTML(value='')))"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Ep 8: loss_policy=0.425\n",
      "Ep 9: reward=105.50\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "c7aee4989d904723bda2390394cf4b29",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(IntProgress(value=0, max=288), HTML(value='')))"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Ep 9: loss_inv=0.000\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "f4fbe96ef7ce454c9eb17c258f492f38",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(IntProgress(value=0, max=32), HTML(value='')))"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Ep 9: loss_policy=0.427\n",
      "Ep 10: reward=86.17\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "90d080b97a354e2486ca8b92cfef7b0d",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(IntProgress(value=0, max=320), HTML(value='')))"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Ep 10: loss_inv=0.000\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "ab2d0394279e40a196471bceef3ffcbf",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(IntProgress(value=0, max=32), HTML(value='')))"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Ep 10: loss_policy=0.397\n",
      "Ep 11: reward=200.00\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "d92b940f13e74e9aa81549c6aa99f949",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(IntProgress(value=0, max=352), HTML(value='')))"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Ep 11: loss_inv=0.000\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "5fff9408a96249b0989e4bbd17780a87",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(IntProgress(value=0, max=32), HTML(value='')))"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Ep 11: loss_policy=0.395\n",
      "Ep 12: reward=200.00\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "304b8bc6fc2e4c099af20e57dbec67fc",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(IntProgress(value=0, max=383), HTML(value='')))"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Ep 12: loss_inv=0.002\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "6e33bd450ff8433a8fd25760bb1107bd",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(IntProgress(value=0, max=32), HTML(value='')))"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Ep 12: loss_policy=0.394\n",
      "Ep 13: reward=200.00\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "7d6f3e9f5ed647e2b55110107dfbbbe2",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(IntProgress(value=0, max=414), HTML(value='')))"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Ep 13: loss_inv=0.000\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "7e250554af9840359d22ac5449875866",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(IntProgress(value=0, max=32), HTML(value='')))"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Ep 13: loss_policy=0.403\n",
      "Ep 14: reward=200.00\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "4c855e8736d348979b37062413c89d5f",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(IntProgress(value=0, max=445), HTML(value='')))"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Ep 14: loss_inv=0.000\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "0c10304c081a4e708ae493cbe022b4f6",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(IntProgress(value=0, max=32), HTML(value='')))"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Ep 14: loss_policy=0.451\n",
      "Ep 15: reward=196.83\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "f73c2aae698b474f9f5ac6e5c860e7f2",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(IntProgress(value=0, max=482), HTML(value='')))"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Ep 15: loss_inv=0.000\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "a21b586c06a24c758e8216eb7c35e0c8",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(IntProgress(value=0, max=32), HTML(value='')))"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Ep 15: loss_policy=0.387\n",
      "\n"
     ]
    }
   ],
   "source": [
    "trajs_inv = []\n",
    "\n",
    "for e in tqdm(range(EPOCHS)):\n",
    "    \n",
    "    # step1, generate inverse samples\n",
    "    cnt = 0\n",
    "    epn = 0\n",
    "    \n",
    "    rews = 0\n",
    "        \n",
    "    while True:\n",
    "        traj = []\n",
    "        rew = 0\n",
    "            \n",
    "        obs = env.reset()\n",
    "        while True:\n",
    "            inp = T.from_numpy(obs).view(((1, )+obs.shape)).float().cuda()\n",
    "            out = model.pred_act(inp).cpu().detach().numpy()\n",
    "                \n",
    "            if np.random.rand()>=EPS:\n",
    "                act = np.argmax(out, axis=1)[0]\n",
    "            else:\n",
    "                act = env.action_space.sample()\n",
    "                \n",
    "            new_obs, r, done, _ = env.step(act)\n",
    "                \n",
    "            traj.append([obs, act, new_obs])\n",
    "            obs = new_obs\n",
    "            rew += r\n",
    "            \n",
    "            cnt += 1\n",
    "                \n",
    "            if done==True:\n",
    "                rews += rew\n",
    "                trajs_inv.append(traj)\n",
    "                \n",
    "                epn += 1\n",
    "                \n",
    "                break\n",
    "        \n",
    "        if cnt >= M:\n",
    "            break\n",
    "        \n",
    "    rews /= epn\n",
    "    print('Ep %d: reward=%.2f' % (e+1, rews))\n",
    "        \n",
    "    # step2, update inverse model\n",
    "    ld_inv = DataLoader(DS_Inv(trajs_inv), batch_size=32, shuffle=True)\n",
    "    \n",
    "    with tqdm(ld_inv) as TQ:\n",
    "        ls_ep = 0\n",
    "        \n",
    "        for obs1, obs2, act in TQ:\n",
    "            out = model.pred_inv(obs1.float().cuda(), obs2.float().cuda())\n",
    "            ls_bh = loss_func(out, act.cuda())\n",
    "            \n",
    "            optim.zero_grad()\n",
    "            ls_bh.backward()\n",
    "            optim.step()\n",
    "            \n",
    "            ls_bh = ls_bh.cpu().detach().numpy()\n",
    "            TQ.set_postfix(loss_inv='%.3f' % (ls_bh))\n",
    "            ls_ep += ls_bh\n",
    "        \n",
    "        ls_ep /= len(TQ)\n",
    "        print('Ep %d: loss_inv=%.3f' % (e+1, ls_ep))\n",
    "    \n",
    "    # step3, predict inverse action for demo samples\n",
    "    traj_policy = []\n",
    "    \n",
    "    for obs1, obs2, _ in ld_demo:\n",
    "        out = model.pred_inv(obs1.float().cuda(), obs2.float().cuda())\n",
    "        \n",
    "        obs = obs1.cpu().detach().numpy()\n",
    "        out = out.cpu().detach().numpy()\n",
    "        out = np.argmax(out, axis=1)\n",
    "        \n",
    "        for i in range(100):\n",
    "            traj_policy.append([obs[i], out[i]])\n",
    "    \n",
    "    # step4, update policy via demo samples\n",
    "    ld_policy = DataLoader(DS_Policy(traj_policy), batch_size=32, shuffle=True)\n",
    "    \n",
    "    with tqdm(ld_policy) as TQ:\n",
    "        ls_ep = 0\n",
    "        \n",
    "        for obs, act in TQ:\n",
    "            out = model.pred_act(obs.float().cuda())\n",
    "            ls_bh = loss_func(out, act.cuda())\n",
    "            \n",
    "            optim.zero_grad()\n",
    "            ls_bh.backward()\n",
    "            optim.step()\n",
    "            \n",
    "            ls_bh = ls_bh.cpu().detach().numpy()\n",
    "            TQ.set_postfix(loss_policy='%.3f' % (ls_bh))\n",
    "            ls_ep += ls_bh\n",
    "        \n",
    "        ls_ep /= len(TQ)\n",
    "        print('Ep %d: loss_policy=%.3f' % (e+1, ls_ep))\n",
    "    \n",
    "    # step5, save model\n",
    "    T.save(model.state_dict(), 'Model/model_cart-pole_%d.pt' % (e+1))\n",
    "    \n",
    "    EPS *= DECAY"
   ]
  }
 ],
 "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.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}