{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 2,
   "source": [
    "%cd learners"
   ],
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "[Errno 2] No such file or directory: 'learners'\n",
      "/home/buehrle/dev/InteractionImitation/scratch/etienne/pillbox/learners\n"
     ]
    }
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "import gym\n",
    "from tqdm import tqdm\n",
    "\n",
    "def rollout(pi, max_steps=1000):\n",
    "    env = gym.make('intersim:intersim-v0')\n",
    "    env.reset() # obs = env.reset()\n",
    "    obs, _, done, _ = env.step(0 * env.action_space.sample())\n",
    "    \n",
    "    _except = lambda o, i: torch.cat((o[:i], o[i+1:]))\n",
    "    \n",
    "    _relative_state_v = lambda obs: torch.stack((\n",
    "        obs[..., 0],\n",
    "        obs[..., 1],\n",
    "        (obs[..., 2]**2 + obs[..., 3]**2).sqrt(),\n",
    "        obs[..., 4],\n",
    "        obs[..., 5],\n",
    "    ), -1)\n",
    "    \n",
    "    for _ in tqdm(range(max_steps)):\n",
    "        pi_obs = [\n",
    "            torch.cat((e.unsqueeze(0), _relative_state_v(_except(o, i)))).unsqueeze(0)\n",
    "            for i, (e, o) in enumerate(zip(obs['state'], obs['relative_state']))\n",
    "        ]\n",
    "        \n",
    "        actions = [pi(o).squeeze() for o in pi_obs]\n",
    "        actions = torch.stack(actions).unsqueeze(1)\n",
    "        obs, _, done, _ = env.step(actions)\n",
    "        env.render(mode='post')\n",
    "        if done:\n",
    "            break\n",
    "    env.close()"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "source": [
    "import torch\n",
    "import numpy as np\n",
    "\n",
    "def expert(obs):\n",
    "    ego = obs[:, 0]\n",
    "    rel = obs[:, 1:]\n",
    "    front = torch.stack((torch.cos(ego[:, 3]), torch.sin(ego[:, 3])), -1)\n",
    "    left = torch.stack((-torch.sin(ego[:, 3]), torch.cos(ego[:, 3])), -1)\n",
    "    df = (rel[:, :, :2] * front.unsqueeze(1)).sum(-1)\n",
    "    dl = (rel[:, :, :2] * left.unsqueeze(1)).sum(-1)\n",
    "\n",
    "    df = torch.where(df.isnan(), np.inf * torch.ones_like(df), df)\n",
    "    dl = torch.where(dl.isnan(), np.inf * torch.ones_like(dl), dl)\n",
    "    rel = torch.where(rel.isnan(), np.inf * torch.ones_like(rel), rel)\n",
    "\n",
    "    # relative speed in direction of position difference vector\n",
    "    vrel = rel[:, :, 2] * (rel[:, :, :2] * torch.stack((\n",
    "        torch.cos(ego[:, 3].unsqueeze(1) + rel[:, :, 3]),\n",
    "        torch.sin(ego[:, 3].unsqueeze(1) + rel[:, :, 3])),\n",
    "        -1)).sum(-1)\n",
    "    vrel = torch.where(vrel.isnan(), np.inf * torch.ones_like(vrel), vrel)\n",
    "    vrel = torch.maximum(vrel, torch.zeros_like(vrel))\n",
    "    \n",
    "    alpha = torch.atan2(dl, df)\n",
    "    d = (rel[:, :, :2] ** 2).sum(-1)\n",
    "    attn = torch.exp(-torch.where(alpha > 0, 0.8*alpha, 1*alpha)**2 - 0.01 * d - 0.1*vrel) \n",
    "    \n",
    "    act = 10 - ego[:, 2] - 20 * attn.sum(-1)\n",
    "    \n",
    "    return act"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "source": [
    "rollout(expert, max_steps=500)"
   ],
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "Vehicle Trajectory Paths: /home/buehrle/dev/InteractionImitation/InteractionSimulator/datasets/trackfiles/DR_USA_Roundabout_FT/vehicle_tracks_000.csv\n",
      "Map Path: /home/buehrle/dev/InteractionImitation/InteractionSimulator/datasets/maps/DR_USA_Roundabout_FT.osm\n",
      "Environment Reset\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stderr",
     "text": [
      "  0%|          | 0/500 [00:00<?, ?it/s]\n"
     ]
    },
    {
     "output_type": "error",
     "ename": "RuntimeError",
     "evalue": "torch.cat(): Sizes of tensors must match except in dimension 0. Got 5 and 6 in dimension 1 (The offending index is 1)",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mRuntimeError\u001b[0m                              Traceback (most recent call last)",
      "\u001b[0;32m/tmp/ipykernel_4266/633354333.py\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[0;32m----> 1\u001b[0;31m \u001b[0mrollout\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mexpert\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mmax_steps\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;36m500\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m",
      "\u001b[0;32m/tmp/ipykernel_4266/1839273282.py\u001b[0m in \u001b[0;36mrollout\u001b[0;34m(pi, max_steps)\u001b[0m\n\u001b[1;32m     12\u001b[0m         pi_obs = [\n\u001b[1;32m     13\u001b[0m             \u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcat\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0me\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0munsqueeze\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;36m0\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0m_except_self\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mo\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mi\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0munsqueeze\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[0m\n\u001b[0;32m---> 14\u001b[0;31m             \u001b[0;32mfor\u001b[0m \u001b[0mi\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m(\u001b[0m\u001b[0me\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mo\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;32min\u001b[0m \u001b[0menumerate\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mzip\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mobs\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m'state'\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mobs\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m'relative_state'\u001b[0m\u001b[0;34m]\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[0m\u001b[1;32m     15\u001b[0m         ]\n\u001b[1;32m     16\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/tmp/ipykernel_4266/1839273282.py\u001b[0m in \u001b[0;36m<listcomp>\u001b[0;34m(.0)\u001b[0m\n\u001b[1;32m     12\u001b[0m         pi_obs = [\n\u001b[1;32m     13\u001b[0m             \u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcat\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0me\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0munsqueeze\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;36m0\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0m_except_self\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mo\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mi\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0munsqueeze\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[0m\n\u001b[0;32m---> 14\u001b[0;31m             \u001b[0;32mfor\u001b[0m \u001b[0mi\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m(\u001b[0m\u001b[0me\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mo\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;32min\u001b[0m \u001b[0menumerate\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mzip\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mobs\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m'state'\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mobs\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m'relative_state'\u001b[0m\u001b[0;34m]\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[0m\u001b[1;32m     15\u001b[0m         ]\n\u001b[1;32m     16\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mRuntimeError\u001b[0m: torch.cat(): Sizes of tensors must match except in dimension 0. Got 5 and 6 in dimension 1 (The offending index is 1)"
     ]
    }
   ],
   "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": 5
}