linear_regression_shuffled.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "import torch.nn as nn\n",
    "from torch.autograd import Variable\n",
    "from sklearn.model_selection import train_test_split\n",
    "import os, os.path \n",
    "import numpy \n",
    "import pickle\n",
    "from glob import glob\n",
    "\n",
    "\"\"\"Change to the data folder\"\"\"\n",
    "train_path = \"./new_train/new_train/\"\n",
    "vali_path = \"./new_val_in/new_val_in\"\n",
    "# number of sequences in each dataset\n",
    "# train:205942  val:3200 test: 36272 \n",
    "# sequences sampled at 10HZ rate"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Create a dataset class "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "class ArgoverseDataset(Dataset):\n",
    "    \"\"\"Dataset class for Argoverse\"\"\"\n",
    "    def __init__(self, data_path: str, transform=None):\n",
    "        super(ArgoverseDataset, self).__init__()\n",
    "        self.data_path = data_path\n",
    "        self.transform = transform\n",
    "\n",
    "        self.pkl_list = glob(os.path.join(self.data_path, '*'))\n",
    "        self.pkl_list.sort()\n",
    "        \n",
    "    def __len__(self):\n",
    "        return len(self.pkl_list)\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "\n",
    "        pkl_path = self.pkl_list[idx]\n",
    "        with open(pkl_path, 'rb') as f:\n",
    "            data = pickle.load(f)\n",
    "            \n",
    "        if self.transform:\n",
    "            data = self.transform(data)\n",
    "\n",
    "        return data\n",
    "\n",
    "\n",
    "# intialize a dataset\n",
    "train_dataset  = ArgoverseDataset(data_path=train_path)\n",
    "vali_dataset = ArgoverseDataset(data_path=vali_path)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Create a loader to enable batch processing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "batch_sz = 4\n",
    "\n",
    "def my_collate(batch):\n",
    "    \"\"\" collate lists of samples into batches, create [ batch_sz x agent_sz x seq_len x feature] \"\"\"\n",
    "    inp = [numpy.dstack([scene['p_in'], scene['v_in']]) for scene in batch]\n",
    "    out = [numpy.dstack([scene['p_out'], scene['v_out']]) for scene in batch]\n",
    "    inp = torch.LongTensor(inp)\n",
    "    out = torch.LongTensor(out)\n",
    "    return [inp, out]\n",
    "\n",
    "train_loader = DataLoader(train_dataset,batch_size=batch_sz, shuffle = False, collate_fn=my_collate, num_workers=0)\n",
    "vali_loader = DataLoader(vali_dataset,batch_size=batch_sz, shuffle = False, collate_fn=my_collate, num_workers=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "205942"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(train_dataset)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "3200"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(vali_dataset)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from sklearn.metrics import mean_squared_error\n",
    "from sklearn.linear_model import LinearRegression"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "#for i in range(0, len(train_dataset)):\n",
    "\n",
    "x = train_dataset[0]['p_in']\n",
    "#    print(type(x))\n",
    "    #print(x.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "import random\n",
    "index = random.sample(range(1, len(train_dataset)), 10000)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "10000"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(index)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "numpy.ndarray"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "type(x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[3277.29638672, 1947.62609863],\n",
       "       [3277.29614258, 1947.62597656],\n",
       "       [3277.29614258, 1947.62585449],\n",
       "       [3277.29614258, 1947.62597656],\n",
       "       [3277.29589844, 1947.62561035],\n",
       "       [3277.29589844, 1947.62561035],\n",
       "       [3277.29589844, 1947.62548828],\n",
       "       [3277.29614258, 1947.62573242],\n",
       "       [3277.29638672, 1947.62573242],\n",
       "       [3277.29614258, 1947.62573242],\n",
       "       [3277.29614258, 1947.62561035],\n",
       "       [3277.29638672, 1947.62573242],\n",
       "       [3277.29638672, 1947.62585449],\n",
       "       [3277.29663086, 1947.62597656],\n",
       "       [3277.29638672, 1947.62585449],\n",
       "       [3277.29638672, 1947.62585449],\n",
       "       [3277.29638672, 1947.62561035],\n",
       "       [3277.29638672, 1947.62573242],\n",
       "       [3277.29638672, 1947.62585449]])"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "n is 0\n",
      "n is 1000\n",
      "n is 2000\n",
      "n is 3000\n",
      "n is 4000\n",
      "n is 5000\n",
      "n is 6000\n",
      "n is 7000\n",
      "n is 8000\n",
      "n is 9000\n"
     ]
    }
   ],
   "source": [
    "p_x = []\n",
    "p_y = []\n",
    "out_x = []\n",
    "out_y = []\n",
    "\n",
    "for n in range(0, 10000):\n",
    "    x = train_dataset[index[n]]['track_id'] \n",
    "    id = train_dataset[index[n]]['agent_id']\n",
    "    i = np.where(x == id)[0][0]\n",
    "    position = train_dataset[index[n]]['p_in'][i]\n",
    "    position = position.reshape(38)\n",
    "    label = train_dataset[index[n]]['p_out'][i]\n",
    "    label = label.reshape(60)\n",
    "    p_x =  p_x + list(position[::2])\n",
    "    p_y = p_y + list(position[1::2])\n",
    "    out_x = out_x + list(label[::2])\n",
    "    out_y = out_y + list(label[1::2])\n",
    "    if(n % 1000 == 0):\n",
    "        print(\"n is \" + str(n))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "190000\n",
      "190000\n",
      "300000\n",
      "300000\n"
     ]
    }
   ],
   "source": [
    "print(len(p_x))\n",
    "print(len(p_y))\n",
    "print(len(out_x))\n",
    "print(len(out_y))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "p_x = np.array(p_x).reshape((10000,19))\n",
    "p_y = np.array(p_y).reshape((10000,19))\n",
    "out_x = np.array(out_x).reshape((10000,30))\n",
    "out_y = np.array(out_y).reshape((10000,30))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_out(n, label):\n",
    "    out = []\n",
    "    for i in range(0,len(label)):\n",
    "        out.append(label[i][n])\n",
    "    return out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "y_0 = get_out(0, out_x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "10000"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(y_0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "n is 0\n",
      "n is 1000\n",
      "n is 2000\n",
      "n is 3000\n"
     ]
    }
   ],
   "source": [
    "t_x= []\n",
    "t_y = []\n",
    "for n in range(0, len(vali_dataset)):\n",
    "    x = vali_dataset[n]['track_id'] \n",
    "    id = vali_dataset[n]['agent_id']\n",
    "    i = np.where(x == id)[0][0]\n",
    "    position = vali_dataset[n]['p_in'][i]\n",
    "    position = position.reshape(38)\n",
    "    t_x =  t_x + list(position[::2])\n",
    "    t_y = t_y + list(position[1::2])\n",
    "    if(n % 1000 == 0):\n",
    "        print(\"n is \" + str(n))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "60800"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(t_x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "60800"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(t_y)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "t_x = np.array(t_x).reshape((3200,19))\n",
    "t_y = np.array(t_y).reshape((3200,19))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "#vali_dataset[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [],
   "source": [
    "s_id = [None] * 3200\n",
    "for i in range(0, len(vali_dataset)):\n",
    "    s_id[i] = vali_dataset[i]['scene_idx']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>0</th>\n",
       "      <th>1</th>\n",
       "      <th>2</th>\n",
       "      <th>3</th>\n",
       "      <th>4</th>\n",
       "      <th>5</th>\n",
       "      <th>6</th>\n",
       "      <th>7</th>\n",
       "      <th>8</th>\n",
       "      <th>9</th>\n",
       "      <th>...</th>\n",
       "      <th>50</th>\n",
       "      <th>51</th>\n",
       "      <th>52</th>\n",
       "      <th>53</th>\n",
       "      <th>54</th>\n",
       "      <th>55</th>\n",
       "      <th>56</th>\n",
       "      <th>57</th>\n",
       "      <th>58</th>\n",
       "      <th>59</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>10002</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>...</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>10015</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>...</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>10019</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>...</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>10028</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>...</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>1003</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>...</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>...</th>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3195</th>\n",
       "      <td>9897</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>...</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3196</th>\n",
       "      <td>99</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>...</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3197</th>\n",
       "      <td>9905</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>...</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3198</th>\n",
       "      <td>9910</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>...</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3199</th>\n",
       "      <td>9918</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>...</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>3200 rows × 60 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "         0    1    2    3    4    5    6    7    8    9   ...   50   51   52  \\\n",
       "0     10002  NaN  NaN  NaN  NaN  NaN  NaN  NaN  NaN  NaN  ...  NaN  NaN  NaN   \n",
       "1     10015  NaN  NaN  NaN  NaN  NaN  NaN  NaN  NaN  NaN  ...  NaN  NaN  NaN   \n",
       "2     10019  NaN  NaN  NaN  NaN  NaN  NaN  NaN  NaN  NaN  ...  NaN  NaN  NaN   \n",
       "3     10028  NaN  NaN  NaN  NaN  NaN  NaN  NaN  NaN  NaN  ...  NaN  NaN  NaN   \n",
       "4      1003  NaN  NaN  NaN  NaN  NaN  NaN  NaN  NaN  NaN  ...  NaN  NaN  NaN   \n",
       "...     ...  ...  ...  ...  ...  ...  ...  ...  ...  ...  ...  ...  ...  ...   \n",
       "3195   9897  NaN  NaN  NaN  NaN  NaN  NaN  NaN  NaN  NaN  ...  NaN  NaN  NaN   \n",
       "3196     99  NaN  NaN  NaN  NaN  NaN  NaN  NaN  NaN  NaN  ...  NaN  NaN  NaN   \n",
       "3197   9905  NaN  NaN  NaN  NaN  NaN  NaN  NaN  NaN  NaN  ...  NaN  NaN  NaN   \n",
       "3198   9910  NaN  NaN  NaN  NaN  NaN  NaN  NaN  NaN  NaN  ...  NaN  NaN  NaN   \n",
       "3199   9918  NaN  NaN  NaN  NaN  NaN  NaN  NaN  NaN  NaN  ...  NaN  NaN  NaN   \n",
       "\n",
       "       53   54   55   56   57   58   59  \n",
       "0     NaN  NaN  NaN  NaN  NaN  NaN  NaN  \n",
       "1     NaN  NaN  NaN  NaN  NaN  NaN  NaN  \n",
       "2     NaN  NaN  NaN  NaN  NaN  NaN  NaN  \n",
       "3     NaN  NaN  NaN  NaN  NaN  NaN  NaN  \n",
       "4     NaN  NaN  NaN  NaN  NaN  NaN  NaN  \n",
       "...   ...  ...  ...  ...  ...  ...  ...  \n",
       "3195  NaN  NaN  NaN  NaN  NaN  NaN  NaN  \n",
       "3196  NaN  NaN  NaN  NaN  NaN  NaN  NaN  \n",
       "3197  NaN  NaN  NaN  NaN  NaN  NaN  NaN  \n",
       "3198  NaN  NaN  NaN  NaN  NaN  NaN  NaN  \n",
       "3199  NaN  NaN  NaN  NaN  NaN  NaN  NaN  \n",
       "\n",
       "[3200 rows x 60 columns]"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dfc = pd.DataFrame(columns=np.arange(60))\n",
    "dfc[0] = list(s_id)\n",
    "dfc"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [],
   "source": [
    "for i in range(0,30):\n",
    "    real_x = get_out(i, out_x)\n",
    "    real_y = get_out(i, out_y)\n",
    "    regressor_x = LinearRegression()\n",
    "    regressor_x.fit(p_x, real_x)\n",
    "    pred_x_i = regressor_x.predict(t_x)\n",
    "    regressor_y = LinearRegression()\n",
    "    regressor_y.fit(p_y, real_y)\n",
    "    pred_y_i = regressor_y.predict(t_y)\n",
    "    dfc[i * 2 + 1] = list(pred_x_i)\n",
    "    dfc[2 * i + 2] = list(pred_y_i)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>0</th>\n",
       "      <th>1</th>\n",
       "      <th>2</th>\n",
       "      <th>3</th>\n",
       "      <th>4</th>\n",
       "      <th>5</th>\n",
       "      <th>6</th>\n",
       "      <th>7</th>\n",
       "      <th>8</th>\n",
       "      <th>9</th>\n",
       "      <th>...</th>\n",
       "      <th>51</th>\n",
       "      <th>52</th>\n",
       "      <th>53</th>\n",
       "      <th>54</th>\n",
       "      <th>55</th>\n",
       "      <th>56</th>\n",
       "      <th>57</th>\n",
       "      <th>58</th>\n",
       "      <th>59</th>\n",
       "      <th>60</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>10002</td>\n",
       "      <td>1714.777378</td>\n",
       "      <td>337.199261</td>\n",
       "      <td>1715.730199</td>\n",
       "      <td>338.116298</td>\n",
       "      <td>1716.695437</td>\n",
       "      <td>339.035481</td>\n",
       "      <td>1717.620307</td>\n",
       "      <td>339.941665</td>\n",
       "      <td>1718.546048</td>\n",
       "      <td>...</td>\n",
       "      <td>1738.764473</td>\n",
       "      <td>361.348703</td>\n",
       "      <td>1739.721823</td>\n",
       "      <td>362.272508</td>\n",
       "      <td>1740.650292</td>\n",
       "      <td>363.280774</td>\n",
       "      <td>1741.531718</td>\n",
       "      <td>364.255287</td>\n",
       "      <td>1742.452859</td>\n",
       "      <td>365.210979</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>10015</td>\n",
       "      <td>725.524088</td>\n",
       "      <td>1229.675534</td>\n",
       "      <td>725.464439</td>\n",
       "      <td>1229.442801</td>\n",
       "      <td>725.388316</td>\n",
       "      <td>1229.221879</td>\n",
       "      <td>725.327991</td>\n",
       "      <td>1229.029871</td>\n",
       "      <td>725.259416</td>\n",
       "      <td>...</td>\n",
       "      <td>723.349056</td>\n",
       "      <td>1226.002466</td>\n",
       "      <td>723.242466</td>\n",
       "      <td>1225.900470</td>\n",
       "      <td>723.141930</td>\n",
       "      <td>1225.814839</td>\n",
       "      <td>723.050571</td>\n",
       "      <td>1225.736556</td>\n",
       "      <td>722.949185</td>\n",
       "      <td>1225.642310</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>10019</td>\n",
       "      <td>574.022163</td>\n",
       "      <td>1244.657588</td>\n",
       "      <td>574.226982</td>\n",
       "      <td>1244.582410</td>\n",
       "      <td>574.462319</td>\n",
       "      <td>1244.498686</td>\n",
       "      <td>574.691099</td>\n",
       "      <td>1244.432921</td>\n",
       "      <td>574.912899</td>\n",
       "      <td>...</td>\n",
       "      <td>579.623051</td>\n",
       "      <td>1242.630450</td>\n",
       "      <td>579.846756</td>\n",
       "      <td>1242.555412</td>\n",
       "      <td>580.065052</td>\n",
       "      <td>1242.478833</td>\n",
       "      <td>580.264209</td>\n",
       "      <td>1242.395315</td>\n",
       "      <td>580.465167</td>\n",
       "      <td>1242.324058</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>10028</td>\n",
       "      <td>1691.344232</td>\n",
       "      <td>315.619272</td>\n",
       "      <td>1691.967599</td>\n",
       "      <td>316.176014</td>\n",
       "      <td>1692.595622</td>\n",
       "      <td>316.757616</td>\n",
       "      <td>1693.238591</td>\n",
       "      <td>317.365832</td>\n",
       "      <td>1693.875006</td>\n",
       "      <td>...</td>\n",
       "      <td>1707.256780</td>\n",
       "      <td>330.964330</td>\n",
       "      <td>1707.896469</td>\n",
       "      <td>331.584855</td>\n",
       "      <td>1708.525031</td>\n",
       "      <td>332.221458</td>\n",
       "      <td>1709.132092</td>\n",
       "      <td>332.830139</td>\n",
       "      <td>1709.747009</td>\n",
       "      <td>333.432889</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>1003</td>\n",
       "      <td>2122.794130</td>\n",
       "      <td>677.105355</td>\n",
       "      <td>2121.513165</td>\n",
       "      <td>676.000026</td>\n",
       "      <td>2120.217831</td>\n",
       "      <td>674.929059</td>\n",
       "      <td>2118.966996</td>\n",
       "      <td>673.840681</td>\n",
       "      <td>2117.680518</td>\n",
       "      <td>...</td>\n",
       "      <td>2090.609261</td>\n",
       "      <td>650.292290</td>\n",
       "      <td>2089.356006</td>\n",
       "      <td>649.279360</td>\n",
       "      <td>2088.113987</td>\n",
       "      <td>648.280420</td>\n",
       "      <td>2086.915795</td>\n",
       "      <td>647.223142</td>\n",
       "      <td>2085.676714</td>\n",
       "      <td>646.221800</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>...</th>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3195</th>\n",
       "      <td>9897</td>\n",
       "      <td>256.429996</td>\n",
       "      <td>805.455028</td>\n",
       "      <td>256.676400</td>\n",
       "      <td>805.249830</td>\n",
       "      <td>256.931963</td>\n",
       "      <td>805.037179</td>\n",
       "      <td>257.182086</td>\n",
       "      <td>804.843303</td>\n",
       "      <td>257.436113</td>\n",
       "      <td>...</td>\n",
       "      <td>263.170258</td>\n",
       "      <td>799.872656</td>\n",
       "      <td>263.448220</td>\n",
       "      <td>799.644081</td>\n",
       "      <td>263.710266</td>\n",
       "      <td>799.437811</td>\n",
       "      <td>263.979462</td>\n",
       "      <td>799.200071</td>\n",
       "      <td>264.243175</td>\n",
       "      <td>798.994559</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3196</th>\n",
       "      <td>99</td>\n",
       "      <td>587.866477</td>\n",
       "      <td>1153.969822</td>\n",
       "      <td>587.897815</td>\n",
       "      <td>1153.366877</td>\n",
       "      <td>587.930806</td>\n",
       "      <td>1152.699491</td>\n",
       "      <td>587.967273</td>\n",
       "      <td>1152.034222</td>\n",
       "      <td>587.997496</td>\n",
       "      <td>...</td>\n",
       "      <td>588.688828</td>\n",
       "      <td>1137.105300</td>\n",
       "      <td>588.718455</td>\n",
       "      <td>1136.419453</td>\n",
       "      <td>588.742999</td>\n",
       "      <td>1135.763588</td>\n",
       "      <td>588.769687</td>\n",
       "      <td>1135.095689</td>\n",
       "      <td>588.801478</td>\n",
       "      <td>1134.462162</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3197</th>\n",
       "      <td>9905</td>\n",
       "      <td>1755.594999</td>\n",
       "      <td>444.202283</td>\n",
       "      <td>1755.428171</td>\n",
       "      <td>444.516914</td>\n",
       "      <td>1755.261726</td>\n",
       "      <td>444.825848</td>\n",
       "      <td>1755.100360</td>\n",
       "      <td>445.139414</td>\n",
       "      <td>1754.932745</td>\n",
       "      <td>...</td>\n",
       "      <td>1751.262173</td>\n",
       "      <td>452.115668</td>\n",
       "      <td>1751.089948</td>\n",
       "      <td>452.423299</td>\n",
       "      <td>1750.920581</td>\n",
       "      <td>452.768510</td>\n",
       "      <td>1750.744512</td>\n",
       "      <td>453.053348</td>\n",
       "      <td>1750.582500</td>\n",
       "      <td>453.379344</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3198</th>\n",
       "      <td>9910</td>\n",
       "      <td>574.607937</td>\n",
       "      <td>1288.846472</td>\n",
       "      <td>574.508950</td>\n",
       "      <td>1288.612320</td>\n",
       "      <td>574.411812</td>\n",
       "      <td>1288.400916</td>\n",
       "      <td>574.316125</td>\n",
       "      <td>1288.202646</td>\n",
       "      <td>574.211579</td>\n",
       "      <td>...</td>\n",
       "      <td>571.628342</td>\n",
       "      <td>1283.345537</td>\n",
       "      <td>571.496566</td>\n",
       "      <td>1283.132435</td>\n",
       "      <td>571.364495</td>\n",
       "      <td>1282.929413</td>\n",
       "      <td>571.230086</td>\n",
       "      <td>1282.702028</td>\n",
       "      <td>571.108232</td>\n",
       "      <td>1282.503128</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3199</th>\n",
       "      <td>9918</td>\n",
       "      <td>584.706860</td>\n",
       "      <td>1164.678087</td>\n",
       "      <td>584.762487</td>\n",
       "      <td>1163.660324</td>\n",
       "      <td>584.812258</td>\n",
       "      <td>1162.684425</td>\n",
       "      <td>584.863561</td>\n",
       "      <td>1161.751364</td>\n",
       "      <td>584.914440</td>\n",
       "      <td>...</td>\n",
       "      <td>585.580341</td>\n",
       "      <td>1141.829012</td>\n",
       "      <td>585.595260</td>\n",
       "      <td>1141.000016</td>\n",
       "      <td>585.611205</td>\n",
       "      <td>1140.133807</td>\n",
       "      <td>585.618539</td>\n",
       "      <td>1139.335423</td>\n",
       "      <td>585.631301</td>\n",
       "      <td>1138.527570</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>3200 rows × 61 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "         0            1            2            3            4            5   \\\n",
       "0     10002  1714.777378   337.199261  1715.730199   338.116298  1716.695437   \n",
       "1     10015   725.524088  1229.675534   725.464439  1229.442801   725.388316   \n",
       "2     10019   574.022163  1244.657588   574.226982  1244.582410   574.462319   \n",
       "3     10028  1691.344232   315.619272  1691.967599   316.176014  1692.595622   \n",
       "4      1003  2122.794130   677.105355  2121.513165   676.000026  2120.217831   \n",
       "...     ...          ...          ...          ...          ...          ...   \n",
       "3195   9897   256.429996   805.455028   256.676400   805.249830   256.931963   \n",
       "3196     99   587.866477  1153.969822   587.897815  1153.366877   587.930806   \n",
       "3197   9905  1755.594999   444.202283  1755.428171   444.516914  1755.261726   \n",
       "3198   9910   574.607937  1288.846472   574.508950  1288.612320   574.411812   \n",
       "3199   9918   584.706860  1164.678087   584.762487  1163.660324   584.812258   \n",
       "\n",
       "               6            7            8            9   ...           51  \\\n",
       "0      339.035481  1717.620307   339.941665  1718.546048  ...  1738.764473   \n",
       "1     1229.221879   725.327991  1229.029871   725.259416  ...   723.349056   \n",
       "2     1244.498686   574.691099  1244.432921   574.912899  ...   579.623051   \n",
       "3      316.757616  1693.238591   317.365832  1693.875006  ...  1707.256780   \n",
       "4      674.929059  2118.966996   673.840681  2117.680518  ...  2090.609261   \n",
       "...           ...          ...          ...          ...  ...          ...   \n",
       "3195   805.037179   257.182086   804.843303   257.436113  ...   263.170258   \n",
       "3196  1152.699491   587.967273  1152.034222   587.997496  ...   588.688828   \n",
       "3197   444.825848  1755.100360   445.139414  1754.932745  ...  1751.262173   \n",
       "3198  1288.400916   574.316125  1288.202646   574.211579  ...   571.628342   \n",
       "3199  1162.684425   584.863561  1161.751364   584.914440  ...   585.580341   \n",
       "\n",
       "               52           53           54           55           56  \\\n",
       "0      361.348703  1739.721823   362.272508  1740.650292   363.280774   \n",
       "1     1226.002466   723.242466  1225.900470   723.141930  1225.814839   \n",
       "2     1242.630450   579.846756  1242.555412   580.065052  1242.478833   \n",
       "3      330.964330  1707.896469   331.584855  1708.525031   332.221458   \n",
       "4      650.292290  2089.356006   649.279360  2088.113987   648.280420   \n",
       "...           ...          ...          ...          ...          ...   \n",
       "3195   799.872656   263.448220   799.644081   263.710266   799.437811   \n",
       "3196  1137.105300   588.718455  1136.419453   588.742999  1135.763588   \n",
       "3197   452.115668  1751.089948   452.423299  1750.920581   452.768510   \n",
       "3198  1283.345537   571.496566  1283.132435   571.364495  1282.929413   \n",
       "3199  1141.829012   585.595260  1141.000016   585.611205  1140.133807   \n",
       "\n",
       "               57           58           59           60  \n",
       "0     1741.531718   364.255287  1742.452859   365.210979  \n",
       "1      723.050571  1225.736556   722.949185  1225.642310  \n",
       "2      580.264209  1242.395315   580.465167  1242.324058  \n",
       "3     1709.132092   332.830139  1709.747009   333.432889  \n",
       "4     2086.915795   647.223142  2085.676714   646.221800  \n",
       "...           ...          ...          ...          ...  \n",
       "3195   263.979462   799.200071   264.243175   798.994559  \n",
       "3196   588.769687  1135.095689   588.801478  1134.462162  \n",
       "3197  1750.744512   453.053348  1750.582500   453.379344  \n",
       "3198   571.230086  1282.702028   571.108232  1282.503128  \n",
       "3199   585.618539  1139.335423   585.631301  1138.527570  \n",
       "\n",
       "[3200 rows x 61 columns]"
      ]
     },
     "execution_count": 26,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dfc"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [],
   "source": [
    "dfc.to_csv('sub_linear_shuffle.csv', index=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [
    {
     "ename": "NameError",
     "evalue": "name 'df' is not defined",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mNameError\u001b[0m                                 Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-28-00cf07b74dcd>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[0;32m----> 1\u001b[0;31m \u001b[0mdf\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m",
      "\u001b[0;31mNameError\u001b[0m: name 'df' is not defined"
     ]
    }
   ],
   "source": [
    "\n",
    "df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#train_dataset[0]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Visualize the batch of sequences"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import random\n",
    "\n",
    "agent_id = 0\n",
    "\n",
    "def show_sample_batch(sample_batch, agent_id):\n",
    "    \"\"\"visualize the trajectory for a batch of samples with a randon agent\"\"\"\n",
    "    inp, out = sample_batch\n",
    "    batch_sz = inp.size(0)\n",
    "    agent_sz = inp.size(1)\n",
    "    \n",
    "    fig, axs = plt.subplots(1,batch_sz, figsize=(15, 3), facecolor='w', edgecolor='k')\n",
    "    fig.subplots_adjust(hspace = .5, wspace=.001)\n",
    "    axs = axs.ravel()   \n",
    "    for i in range(batch_sz):\n",
    "        axs[i].xaxis.set_ticks([])\n",
    "        axs[i].yaxis.set_ticks([])\n",
    "        \n",
    "        # first two feature dimensions are (x,y) positions\n",
    "        axs[i].scatter(inp[i, agent_id,:,0], inp[i, agent_id,:,1])\n",
    "        axs[i].scatter(out[i, agent_id,:,0], out[i, agent_id,:,1])\n",
    "\n",
    "        \n",
    "for i_batch, sample_batch in enumerate(train_loader):\n",
    "    inp, out = sample_batch\n",
    "    \"\"\"TODO:\n",
    "      Deep learning model\n",
    "      training routine\n",
    "    \"\"\"\n",
    "    show_sample_batch(sample_batch, agent_id)\n",
    "    \n",
    "    break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "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.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}