{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 2,
   "source": [
    "%cd learners"
   ],
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "/home/buehrle/dev/InteractionImitation/scratch/etienne/pillbox/learners\n"
     ]
    }
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "source": [
    "import torch\n",
    "import numpy as np"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "source": [
    "# save expert demos to ../experts/Intersim/demos.npz\n",
    "# make sure to split different experts up\n",
    "\n",
    "from intersim.envs.simulator import InteractionSimulator\n",
    "from intersim.utils import get_map_path, get_svt, SVT_to_stateactions\n",
    "import gym\n",
    "from tqdm import tqdm\n",
    "\n",
    "def pillbox_demo(observations, actions, rewards):\n",
    "    demo = {\n",
    "        'env': 'intersim:intersim-v0',\n",
    "        'num_trajs': len(observations),\n",
    "        'mean_reward': rewards.mean(),\n",
    "        'std_reward': rewards.std(),\n",
    "    }\n",
    "    demo.update({\n",
    "        str(i): {\n",
    "            'states': o,\n",
    "            'actions': a,\n",
    "        } for i, (o, a) in enumerate(zip(observations, actions))\n",
    "    })\n",
    "    return demo\n",
    "\n",
    "def intersim_expert_demos(loc, track):\n",
    "    svt, svt_path = get_svt(loc, track)\n",
    "    osm = get_map_path(loc)\n",
    "    \n",
    "    n_actors = svt.simstate.size(1)\n",
    "    observations = []\n",
    "    actions = [] ##\n",
    "    #states, actions = SVT_to_stateactions(svt) ##\n",
    "    rewards = []\n",
    "    \n",
    "    print('Simulating')\n",
    "    env = gym.make('intersim:intersim-v0', svt=svt, map_path=osm)\n",
    "    obs, info = env.reset()\n",
    "    for s in tqdm(svt.simstate[1:]): ##\n",
    "    #for a in actions: ##\n",
    "        relative_state = torch.stack((\n",
    "            obs['relative_state'][..., 0],\n",
    "            obs['relative_state'][..., 1],\n",
    "            (obs['relative_state'][..., 2]**2 + obs['relative_state'][..., 3]**2).sqrt(),\n",
    "            obs['relative_state'][..., 4],\n",
    "            obs['relative_state'][..., 5],\n",
    "        ), -1)\n",
    "        observations.append(torch.cat((\n",
    "            obs['state'].unsqueeze(1),\n",
    "            relative_state,\n",
    "        ), 1))\n",
    "        obs, r, done, info = env.step(env.target_state(s, mu=.01))\n",
    "        #obs, r, done, info = env.step(a) ##\n",
    "        actions.append(info['action_taken'])\n",
    "        rewards.append(r)\n",
    "        assert not done, 'Episode terminated during expert demonstration.'\n",
    "\n",
    "    _except_idx = lambda o, i: torch.cat((o[:i], o[i+1:]))\n",
    "        \n",
    "    # transpose to per-agent observations and actions\n",
    "    print('Transposing')\n",
    "    observations = [torch.stack([_except_idx(o[i], i+1) for o in observations]) for i in range(n_actors)]\n",
    "    actions = [torch.stack([a[i] for a in actions]) for i in range(n_actors)]\n",
    "    \n",
    "    print('Trimming')\n",
    "    # trim observations and actions to start/end of trajectory\n",
    "    _alive = lambda o: (~o.isnan().all(2).all(1)).nonzero()\n",
    "    _start = lambda o: _alive(o).min()\n",
    "    _end = lambda o: _alive(o).max() + 1\n",
    "    start_end = [(_start(obs), _end(obs)) for obs in observations]\n",
    "    observations = [obs[start:end] for obs, (start, end) in zip(observations, start_end)]\n",
    "    actions = [act[start:end] for act, (start, end) in zip(actions, start_end)]\n",
    "    \n",
    "    #print('Cropping')\n",
    "    ## crop observations to max number of observations\n",
    "    #max_num_obs = max([(~obs.isnan().all(2)).sum(1).max() for obs in observations])\n",
    "    #observations = [obs[:, :max_num_obs] for obs in observations]\n",
    "    \n",
    "    observations = [o.numpy() for o in observations]\n",
    "    actions = [a.numpy() for a in actions]\n",
    "    rewards = np.array(rewards)\n",
    "    \n",
    "    return pillbox_demo(observations, actions, rewards)"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "source": [
    "demos = intersim_expert_demos(loc=0, track=0)"
   ],
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "Simulating\n",
      "Custom Vehicle Trajectory Paths\n",
      "Map Path: datasets/maps/DR_USA_Roundabout_FT.osm\n",
      "Environment Reset\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stderr",
     "text": [
      "100%|███████████████████████████████████████████████████████████████████████████████████████████████████████| 3006/3006 [01:17<00:00, 38.87it/s]\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "Transposing\n",
      "Trimming\n"
     ]
    }
   ],
   "metadata": {
    "scrolled": true,
    "tags": [
     "outputPrepend"
    ]
   }
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "source": [
    "demos['num_trajs']"
   ],
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": [
       "151"
      ]
     },
     "metadata": {},
     "execution_count": 6
    }
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "source": [
    "demos['25']['states'].shape"
   ],
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": [
       "(71, 151, 5)"
      ]
     },
     "metadata": {},
     "execution_count": 7
    }
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "source": [
    "np.savez('../experts/intersim:intersim-v0/demos.npz', **demos)"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [],
   "outputs": [],
   "metadata": {}
  }
 ],
 "metadata": {
  "kernelspec": {
   "name": "python3",
   "display_name": "Python 3.7.5 64-bit ('.venv': venv)"
  },
  "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.5"
  },
  "interpreter": {
   "hash": "56465d2ea10f338edb3d30adb010c5849fd826fffc543ba31360f3db8b47a703"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}