249 lines
6.9 KiB
Plaintext
249 lines
6.9 KiB
Plaintext
{
|
|
"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
|
|
} |