Merge branch 'main' of github.com:sisl/InteractionImitation
This commit is contained in:
Johannes Fischer
@@ -48,6 +48,8 @@ def parse_args():
|
||||
help='seed')
|
||||
parser.add_argument('--nframes', default=500, type=int,
|
||||
help='frames for test animation')
|
||||
parser.add_argument('--nsamples', default=200, type=int,
|
||||
help='number of ray samples')
|
||||
parser.add_argument('--graph', action='store_true',
|
||||
help='whether to mask the relative states based on a ConeVisibilityGraph')
|
||||
parser.add_argument('-d', default='./expert_data', type=str,
|
||||
@@ -64,6 +66,7 @@ def parse_args():
|
||||
'seed':args.seed,
|
||||
'ray':args.ray,
|
||||
'nframes':args.nframes,
|
||||
'nsamples':args.nsamples,
|
||||
'datadir':os.path.abspath(args.d),
|
||||
'graph':None,
|
||||
'outdir': opj('output',args.method,'loc%02i'%(args.loc)),
|
||||
@@ -156,7 +159,7 @@ if __name__ == '__main__':
|
||||
local_dir=kwargs['outdir'],
|
||||
#resources_per_trial={"cpu": 2},
|
||||
time_budget_s=120*60,
|
||||
num_samples=200,
|
||||
num_samples=kwargs['nsamples'],
|
||||
)
|
||||
elif kwargs['ray'] and kwargs['test']:
|
||||
analysis = Analysis(kwargs['outdir'], default_metric="cv_loss", default_mode="min")
|
||||
|
||||
274
scratch/etienne/pillbox/intersim_advil.ipynb
Normal file
274
scratch/etienne/pillbox/intersim_advil.ipynb
Normal file
File diff suppressed because one or more lines are too long
249
scratch/etienne/pillbox/intersim_demos.ipynb
Normal file
249
scratch/etienne/pillbox/intersim_demos.ipynb
Normal file
@@ -0,0 +1,249 @@
|
||||
{
|
||||
"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
|
||||
}
|
||||
168
scratch/etienne/pillbox/intersim_expert.ipynb
Normal file
168
scratch/etienne/pillbox/intersim_expert.ipynb
Normal file
@@ -0,0 +1,168 @@
|
||||
{
|
||||
"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
|
||||
}
|
||||
2121
scratch/etienne/pillbox/intersim_stats.ipynb
Normal file
2121
scratch/etienne/pillbox/intersim_stats.ipynb
Normal file
File diff suppressed because one or more lines are too long
140
scratch/etienne/pillbox/learners/adril.py
Normal file
140
scratch/etienne/pillbox/learners/adril.py
Normal file
@@ -0,0 +1,140 @@
|
||||
import gym
|
||||
from gym import spaces
|
||||
from sklearn.neighbors import KDTree
|
||||
from scipy.stats import norm
|
||||
import numpy as np
|
||||
import warnings
|
||||
from abc import ABC, abstractmethod
|
||||
from typing import Dict, Generator, Optional, Union
|
||||
import torch as th
|
||||
|
||||
try:
|
||||
# Check memory used by replay buffer when possible
|
||||
import psutil
|
||||
except ImportError:
|
||||
psutil = None
|
||||
|
||||
from stable_baselines3.common.preprocessing import get_action_dim, get_obs_shape
|
||||
from stable_baselines3.common.type_aliases import ReplayBufferSamples, RolloutBufferSamples
|
||||
from stable_baselines3.common.vec_env import VecNormalize
|
||||
from stable_baselines3.common.buffers import ReplayBuffer
|
||||
|
||||
|
||||
class AdRILWrapper(gym.Env):
|
||||
metadata = {'render.modes': ['human']}
|
||||
|
||||
def __init__(self, base_env):
|
||||
super(AdRILWrapper, self).__init__()
|
||||
self.base_env = base_env
|
||||
self.iter = 0
|
||||
self.observation_space = self.base_env.observation_space
|
||||
self.action_space = self.base_env.action_space
|
||||
self.trajs = list()
|
||||
self.num_trajs = 0
|
||||
self.curr_state = None
|
||||
def step(self, action):
|
||||
next_obs, _, done, info = self.base_env.step(action)
|
||||
reward = self.iter # Transformed by replay buffer
|
||||
self.trajs.append((self.curr_state, action, next_obs, done))
|
||||
if done:
|
||||
self.num_trajs += 1
|
||||
self.curr_state = next_obs
|
||||
return next_obs, reward, done, info
|
||||
def reset(self):
|
||||
obs = self.base_env.reset()
|
||||
self.curr_state = obs
|
||||
return obs
|
||||
def render(self, mode='human'):
|
||||
self.base_env.render(mode=mode)
|
||||
def close (self):
|
||||
self.base_env.close()
|
||||
def get_learner_trajs(self):
|
||||
return self.trajs
|
||||
def set_iter(self, k):
|
||||
self.iter = k
|
||||
|
||||
class AdRILReplayBuffer(ReplayBuffer):
|
||||
def __init__(
|
||||
self,
|
||||
buffer_size: int,
|
||||
observation_space: spaces.Space,
|
||||
action_space: spaces.Space,
|
||||
device: Union[th.device, str] = "cpu",
|
||||
n_envs: int = 1,
|
||||
optimize_memory_usage: bool = False,
|
||||
expert_data: dict = dict(),
|
||||
N_expert: int = 0,
|
||||
balanced: bool = True,
|
||||
):
|
||||
super(AdRILReplayBuffer, self).__init__(buffer_size, observation_space, action_space, device, n_envs=n_envs, optimize_memory_usage=optimize_memory_usage)
|
||||
|
||||
self.expert_states = expert_data['obs']
|
||||
self.expert_actions = expert_data['acts']
|
||||
self.expert_next_states = expert_data['next_obs']
|
||||
self.expert_dones = expert_data['dones']
|
||||
n_expert = len(expert_data["obs"])
|
||||
self.iter = 0
|
||||
self.N_expert = N_expert
|
||||
self.N_learner = 0
|
||||
self.normalizer = 1
|
||||
self.balanced = balanced
|
||||
|
||||
def set_iter(self, k):
|
||||
self.iter = k
|
||||
normalizer = 0
|
||||
for i in range(0, k):
|
||||
normalizer += 1 ** (-i) # written to support decaying learning rate
|
||||
self.normalizer = normalizer
|
||||
|
||||
def set_n_learner(self, n):
|
||||
self.N_learner = n
|
||||
|
||||
def _get_samples(self, batch_inds: np.ndarray, env: Optional[VecNormalize] = None) -> ReplayBufferSamples:
|
||||
num_samples = len(batch_inds)
|
||||
if self.balanced:
|
||||
num_expert_samples = int(num_samples / 2)
|
||||
batch_inds = batch_inds[:num_expert_samples]
|
||||
expert_inds = np.random.randint(0, len(self.expert_states), size=num_expert_samples)
|
||||
# balanced sampling
|
||||
if self.optimize_memory_usage:
|
||||
next_obs = self._normalize_obs(self.observations[(batch_inds + 1) % self.buffer_size, 0, :], env)
|
||||
else:
|
||||
next_obs = self._normalize_obs(self.next_observations[batch_inds, 0, :], env)
|
||||
next_obs = np.concatenate((next_obs, self._normalize_obs(self.expert_next_states[expert_inds], env)), axis=0)
|
||||
obs = self._normalize_obs(self.observations[batch_inds, 0, :], env)
|
||||
obs = np.concatenate((obs, self._normalize_obs(self.expert_states[expert_inds], env)), axis=0)
|
||||
actions = self.actions[batch_inds, 0, :]
|
||||
actions = np.concatenate((actions, self.expert_actions[expert_inds].reshape(num_expert_samples, -1)), axis=0)
|
||||
dones = self.dones[batch_inds]
|
||||
dones = np.concatenate((dones, self.expert_dones[expert_inds].reshape(num_expert_samples, -1)), axis=0)
|
||||
# AdRIL Rewards (indicator kernel)
|
||||
mask1 = (self.rewards[batch_inds] >= 0).astype(np.float32)
|
||||
mask2 = (self.rewards[batch_inds] < self.iter).astype(np.float32)
|
||||
r1 = - (1. ** (-self.rewards[batch_inds])) * mask1 * mask2 # Past iter
|
||||
r2 = np.zeros_like(self.rewards[batch_inds]) * mask1 * (1 - mask2) # current iter
|
||||
r3 = -self.rewards[batch_inds] * (1 - mask1) # Expert
|
||||
if self.iter > 0:
|
||||
rewards = (r1 / self.N_learner) + r2 + r3
|
||||
else:
|
||||
rewards = r1 + r2 + r3
|
||||
rewards = np.concatenate((rewards, np.ones_like(rewards) / self.N_expert), axis=0)
|
||||
else:
|
||||
if self.optimize_memory_usage:
|
||||
next_obs = self._normalize_obs(self.observations[(batch_inds + 1) % self.buffer_size, 0, :], env)
|
||||
else:
|
||||
next_obs = self._normalize_obs(self.next_observations[batch_inds, 0, :], env)
|
||||
obs = self._normalize_obs(self.observations[batch_inds, 0, :], env)
|
||||
actions = self.actions[batch_inds, 0, :]
|
||||
dones = self.dones[batch_inds]
|
||||
# AdRIL Rewards (indicator kernel)
|
||||
mask1 = (self.rewards[batch_inds] >= 0).astype(np.float32)
|
||||
mask2 = (self.rewards[batch_inds] < self.iter).astype(np.float32)
|
||||
r1 = - (1. ** (-self.rewards[batch_inds])) * mask1 * mask2 # Past iter
|
||||
r2 = np.zeros_like(self.rewards[batch_inds]) * mask1 * (1 - mask2) # current iter
|
||||
r3 = -self.rewards[batch_inds] * (1 - mask1) / self.N_expert # Expert
|
||||
if self.iter > 0:
|
||||
rewards = (r1 * 1. / self.N_learner) + r2 + r3
|
||||
else:
|
||||
rewards = r1 + r2 + r3
|
||||
data = (obs, actions, next_obs, dones, rewards)
|
||||
return ReplayBufferSamples(*tuple(map(self.to_torch, data)))
|
||||
222
scratch/etienne/pillbox/learners/advil.py
Normal file
222
scratch/etienne/pillbox/learners/advil.py
Normal file
@@ -0,0 +1,222 @@
|
||||
import numpy as np
|
||||
|
||||
import torch
|
||||
import torch.autograd as autograd
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
import torch.optim as optim
|
||||
from gym.spaces import Discrete
|
||||
import gym
|
||||
from stable_baselines3.common.preprocessing import get_action_dim
|
||||
from tqdm import tqdm
|
||||
from torch.autograd import Variable
|
||||
from itertools import repeat
|
||||
from torch.autograd import grad as torch_grad
|
||||
from typing import List, Type
|
||||
import types
|
||||
|
||||
# Infinite dataloader
|
||||
def repeater(data_loader):
|
||||
for loader in repeat(data_loader):
|
||||
for data in loader:
|
||||
yield data
|
||||
|
||||
def create_mlp(
|
||||
input_dim: int, output_dim: int, net_arch: List[int], activation_fn: Type[nn.Module] = nn.ReLU) -> List[nn.Module]:
|
||||
|
||||
if len(net_arch) > 0:
|
||||
modules = [nn.Linear(input_dim, net_arch[0]), activation_fn()]
|
||||
else:
|
||||
modules = []
|
||||
|
||||
for idx in range(len(net_arch) - 1):
|
||||
modules.append(nn.Linear(net_arch[idx], net_arch[idx + 1]))
|
||||
modules.append(activation_fn())
|
||||
|
||||
if output_dim > 0:
|
||||
last_layer_dim = net_arch[-1] if len(net_arch) > 0 else input_dim
|
||||
modules.append(nn.Linear(last_layer_dim, output_dim))
|
||||
return modules
|
||||
|
||||
def init_ortho(layer):
|
||||
if type(layer) == nn.Linear:
|
||||
nn.init.orthogonal_(layer.weight)
|
||||
|
||||
|
||||
class AdVILPolicy(nn.Module):
|
||||
def __init__(self, env, mean=None, std=None):
|
||||
super(AdVILPolicy, self).__init__()
|
||||
if isinstance(env.action_space, Discrete):
|
||||
self.net_arch = [64, 64]
|
||||
self.action_dim = env.action_space.n
|
||||
self.discrete = True
|
||||
else:
|
||||
self.net_arch = [256, 256]
|
||||
self.action_dim = int(np.prod(env.action_space.shape))
|
||||
self.low = torch.as_tensor(env.action_space.low)
|
||||
self.high = torch.as_tensor(env.action_space.high)
|
||||
self.discrete = False
|
||||
self.obs_dim = int(np.prod(env.observation_space.shape))
|
||||
self.observation_space = env.observation_space
|
||||
net = create_mlp(self.obs_dim, self.action_dim, self.net_arch, nn.ReLU)
|
||||
if self.discrete:
|
||||
net.append(nn.Softmax(dim=1))
|
||||
self.net = nn.Sequential(*net)
|
||||
self.net.apply(init_ortho)
|
||||
if mean is not None and std is not None:
|
||||
self.mean = mean
|
||||
self.std = std
|
||||
self.is_normalized = True
|
||||
else:
|
||||
self.is_normalized = False
|
||||
def forward(self, obs):
|
||||
action = self.net(obs)
|
||||
return action
|
||||
def predict(self, obs, state, mask, deterministic):
|
||||
obs = obs.reshape((-1,) + (self.obs_dim,))
|
||||
if self.is_normalized:
|
||||
obs = (obs - self.mean) / self.std
|
||||
obs = torch.as_tensor(obs)
|
||||
with torch.no_grad():
|
||||
actions = self.forward(obs)
|
||||
if self.discrete:
|
||||
actions = actions.argmax(dim=1).reshape(-1)
|
||||
else:
|
||||
actions = self.low + ((actions + 1.0) / 2.0) * (self.high - self.low)
|
||||
actions = torch.max(torch.min(actions, self.high), self.low)
|
||||
actions = actions.cpu().numpy()
|
||||
return actions, state
|
||||
|
||||
|
||||
class AdVILDiscriminator(nn.Module):
|
||||
def __init__(self, env):
|
||||
super(AdVILDiscriminator, self).__init__()
|
||||
if isinstance(env.action_space, Discrete):
|
||||
self.net_arch = [64, 64]
|
||||
self.action_dim = env.action_space.n
|
||||
else:
|
||||
self.net_arch = [256, 256]
|
||||
self.action_dim = int(np.prod(env.action_space.shape))
|
||||
self.obs_dim = int(np.prod(env.observation_space.shape))
|
||||
net = create_mlp(self.obs_dim + self.action_dim, 1, self.net_arch, nn.ReLU)
|
||||
self.net = nn.Sequential(*net)
|
||||
self.net.apply(init_ortho)
|
||||
|
||||
def forward(self, inputs):
|
||||
output = self.net(inputs)
|
||||
return output.view(-1)
|
||||
|
||||
def pi_update(obs, acts, pi, f, pi_opt, prog):
|
||||
pi_opt.zero_grad()
|
||||
obs_v = Variable(obs)
|
||||
pi_acts = pi(obs_v)
|
||||
#learner_sa = torch.cat((obs, pi_acts), axis=1)
|
||||
f_learner = f(obs, acts)
|
||||
pi_loss = f_learner.mean() + orthogonal_reg(pi) + 2e-1 * (pi_acts - acts).square().mean()
|
||||
pi_loss.backward()
|
||||
if prog > 0.1:
|
||||
torch.nn.utils.clip_grad_norm(pi.parameters(), 40.0)
|
||||
pi_opt.step()
|
||||
return pi_loss.item(), (2e-1 * (pi_acts - acts).square().mean()).item()
|
||||
|
||||
def orthogonal_reg(pi):
|
||||
with torch.enable_grad():
|
||||
reg = 1e-4
|
||||
orth_loss = torch.zeros(1)
|
||||
for name, param in pi.named_parameters():
|
||||
if 'bias' not in name:
|
||||
x = torch.mm(torch.t(param), param)
|
||||
x = x * (1. - torch.eye(param.shape[-1]))
|
||||
orth_loss = orth_loss + reg * (x.square().sum())
|
||||
return orth_loss
|
||||
|
||||
def f_update(obs, acts, pi, f, f_opt, prog):
|
||||
obs_v = Variable(obs)
|
||||
pi_acts = pi(obs_v)
|
||||
#learner_sa = torch.cat((obs, pi_acts), axis=1)
|
||||
#expert_sa = Variable(torch.cat((obs, acts), axis=1))
|
||||
f_learner = f(obs, pi_acts)
|
||||
f_expert = f(obs, acts)
|
||||
#gp = gradient_penalty((obs, pi_acts), (obs, acts), f)
|
||||
f_opt.zero_grad()
|
||||
f_loss = f_expert.mean() - f_learner.mean()# + 10 * gp
|
||||
f_loss.backward()
|
||||
if prog > 0.1:
|
||||
torch.nn.utils.clip_grad_norm(f.parameters(), 40.0)
|
||||
f_opt.step()
|
||||
return f_loss.item()
|
||||
|
||||
def gradient_penalty(learner_sa, expert_sa, f):
|
||||
batch_size = expert_sa[0].size()[0]
|
||||
|
||||
#alpha = torch.rand(batch_size, 1)
|
||||
#alpha = alpha.expand_as(expert_sa)
|
||||
|
||||
salpha = torch.rand(batch_size, 1, 1)
|
||||
salpha = salpha.expand_as(expert_sa[0])
|
||||
|
||||
aalpha = torch.rand(batch_size, 1)
|
||||
aalpha = aalpha.expand_as(expert_sa[1])
|
||||
|
||||
#interpolated = alpha * expert_sa.data + (1 - alpha) * learner_sa.data
|
||||
#interpolated = Variable(interpolated, requires_grad=True)
|
||||
#f_interpolated = f(interpolated.float())
|
||||
|
||||
sinterpolated = salpha * expert_sa[0].data + (1 - salpha) * learner_sa[0].data
|
||||
sinterpolated = Variable(sinterpolated, requires_grad=True)
|
||||
|
||||
ainterpolated = aalpha * expert_sa[1].data + (1 - aalpha) * learner_sa[1].data
|
||||
ainterpolated = Variable(ainterpolated, requires_grad=True)
|
||||
|
||||
f_interpolated = f(sinterpolated, ainterpolated)
|
||||
|
||||
#gradients = torch_grad(outputs=f_interpolated, inputs=interpolated,
|
||||
# grad_outputs=torch.ones(f_interpolated.size()),
|
||||
# create_graph=True, retain_graph=True)[0]
|
||||
|
||||
sgradients = torch_grad(outputs=f_interpolated, inputs=sinterpolated,
|
||||
grad_outputs=torch.ones(f_interpolated.size()),
|
||||
create_graph=True, retain_graph=True)[0]
|
||||
|
||||
agradients = torch_grad(outputs=f_interpolated, inputs=ainterpolated,
|
||||
grad_outputs=torch.ones(f_interpolated.size()),
|
||||
create_graph=True, retain_graph=True)[0]
|
||||
|
||||
#gradients = gradients.view(batch_size, -1)
|
||||
sgradients = sgradients.view(batch_size, -1)
|
||||
agradients = agradients.view(batch_size, -1)
|
||||
#norm = gradients.norm(2, dim=1).mean().item()
|
||||
#gradients_norm = torch.sqrt(torch.sum(gradients ** 2, dim=1) + 1e-12)
|
||||
gradients_norm = torch.sqrt(torch.sum(sgradients ** 2, dim=1) + torch.sum(agradients ** 2, dim=1) + 1e-12)
|
||||
# 2 * |f'(x_0)|
|
||||
return ((gradients_norm - 0.4) ** 2).mean()
|
||||
|
||||
def advil_training(data_loader, env, iters=int(1e5), policy_class=AdVILPolicy, discriminator_class=AdVILDiscriminator, lr_pi=8e-6, lr_f=8e-4):
|
||||
if not isinstance(env.action_space, Discrete):
|
||||
low = torch.as_tensor(env.action_space.low)
|
||||
high = torch.as_tensor(env.action_space.high)
|
||||
if data_loader.dataset.is_normalized:
|
||||
pi = policy_class(env, data_loader.dataset.mean, data_loader.dataset.std)
|
||||
else:
|
||||
pi = policy_class(env)
|
||||
f = discriminator_class(env)
|
||||
pi_opt = optim.Adam(pi.parameters(), lr=lr_pi)
|
||||
|
||||
last_loss = 0
|
||||
f_opt = optim.Adam(f.parameters(), lr=lr_f)
|
||||
data_loader = repeater(data_loader)
|
||||
for t in tqdm(range(iters)):
|
||||
data = next(data_loader)
|
||||
obs = data['obs']
|
||||
acts = data['acts']
|
||||
#if isinstance(env.action_space, Discrete):
|
||||
# acts = nn.functional.one_hot(acts, env.action_space.n)
|
||||
#else:
|
||||
# acts = (((acts - low) / (high - low)) * 2.0) - 1.0
|
||||
pi_loss, mse_reg = pi_update(obs, acts, pi, f, pi_opt, t/iters)
|
||||
f_loss = f_update(obs, acts, pi, f, f_opt, t/iters)
|
||||
if t % 100 == 0:
|
||||
print("pi loss:", pi_loss)
|
||||
print("mse reg:", mse_reg)
|
||||
print("f loss:", f_loss)
|
||||
return pi
|
||||
155
scratch/etienne/pillbox/learners/intersim_advil.py
Normal file
155
scratch/etienne/pillbox/learners/intersim_advil.py
Normal file
@@ -0,0 +1,155 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
def unnormalize(val, mean, std):
|
||||
val *= std or 1
|
||||
val += mean or 0
|
||||
return val
|
||||
|
||||
def normalize(val, mean, std):
|
||||
val -= mean or 0
|
||||
val /= std or 1
|
||||
return val
|
||||
|
||||
class IntersimPolicy(nn.Module):
|
||||
def __init__(self, env, mean=None, std=None):
|
||||
# assert "intersim" in env.unwrapped.spec.id
|
||||
super().__init__()
|
||||
|
||||
self._ego_encoder = nn.Sequential(
|
||||
# in 5, out 5
|
||||
nn.Linear(5, 10),
|
||||
nn.ReLU(),
|
||||
nn.Linear(10, 10),
|
||||
nn.ReLU(),
|
||||
nn.Linear(10, 10),
|
||||
nn.ReLU(),
|
||||
nn.Linear(10, 5),
|
||||
nn.ReLU(),
|
||||
)
|
||||
self._state_encoder = nn.Sequential(
|
||||
# in 5, out 5
|
||||
nn.Linear(5, 10),
|
||||
nn.ReLU(),
|
||||
nn.Linear(10, 10),
|
||||
nn.ReLU(),
|
||||
nn.Linear(10, 10),
|
||||
nn.ReLU(),
|
||||
nn.Linear(10, 5),
|
||||
nn.ReLU(),
|
||||
)
|
||||
self._deepset = lambda e: e.sum(-2)
|
||||
self._action_decoder = nn.Sequential(
|
||||
# in 5 + 5, out 1
|
||||
nn.Linear(5 + 5, 10),
|
||||
nn.ReLU(),
|
||||
nn.Linear(10, 10),
|
||||
nn.ReLU(),
|
||||
nn.Linear(10, 10),
|
||||
nn.ReLU(),
|
||||
nn.Linear(10, 10),
|
||||
nn.ReLU(),
|
||||
nn.Linear(10, 1),
|
||||
)
|
||||
|
||||
def forward(self, obs):
|
||||
# obs.shape = (batch=514, 1 + others=150, 5)
|
||||
# act.shape = (batch=514, 1)
|
||||
|
||||
ego = obs[:, 0]#.detach().clone()
|
||||
rel = obs[:, 1:]#.detach().clone()
|
||||
nan = rel.isnan().any(-1, keepdim=True)
|
||||
rel = torch.where(nan, torch.zeros_like(rel), rel) # required because of https://github.com/pytorch/pytorch/issues/15506
|
||||
|
||||
d = (rel[:, :, :2] ** 2).sum(-1).sqrt()
|
||||
front = torch.stack((torch.cos(ego[:, 3]), torch.sin(ego[:, 3])), -1)
|
||||
left = torch.stack((-torch.sin(ego[:, 3]), torch.cos(ego[:, 3])), -1)
|
||||
df = (rel[:, :, :2] * front.unsqueeze(1)).sum(-1)
|
||||
dl = (rel[:, :, :2] * left.unsqueeze(1)).sum(-1)
|
||||
alpha = torch.atan2(dl, df)
|
||||
|
||||
rel[:, :, 0] = d
|
||||
rel[:, :, 1] = alpha
|
||||
|
||||
e = self._ego_encoder(ego)
|
||||
x = self._state_encoder(rel)
|
||||
x = torch.where(nan, torch.zeros_like(x), x)
|
||||
x = self._deepset(x)
|
||||
a = self._action_decoder(torch.cat((e, x), 1))
|
||||
|
||||
return 10 * a
|
||||
|
||||
def predict(self, state, mask, deterministic):
|
||||
#action_distribution = self.forward(obs)
|
||||
#action = action_distribution.argmax()
|
||||
#return action
|
||||
return self.forward(obs)
|
||||
|
||||
class IntersimDiscriminator(nn.Module):
|
||||
def __init__(self, env):
|
||||
# assert "intersim" in env.unwrapped.spec.id
|
||||
super().__init__()
|
||||
|
||||
self._ego_encoder = nn.Sequential(
|
||||
# in 5, out 5
|
||||
nn.Linear(5, 10),
|
||||
nn.ReLU(),
|
||||
nn.Linear(10, 10),
|
||||
nn.ReLU(),
|
||||
nn.Linear(10, 10),
|
||||
nn.ReLU(),
|
||||
nn.Linear(10, 5),
|
||||
nn.ReLU(),
|
||||
)
|
||||
self._state_encoder = nn.Sequential(
|
||||
# in 5, out 5
|
||||
nn.Linear(5, 10),
|
||||
nn.ReLU(),
|
||||
nn.Linear(10, 10),
|
||||
nn.ReLU(),
|
||||
nn.Linear(10, 10),
|
||||
nn.ReLU(),
|
||||
nn.Linear(10, 5),
|
||||
nn.ReLU(),
|
||||
)
|
||||
self._deepset = lambda e: e.sum(-2)
|
||||
self._discriminator = nn.Sequential(
|
||||
# in 5 + 5 + 1, out 1
|
||||
nn.Linear(5 + 5 + 1, 10),
|
||||
nn.ReLU(),
|
||||
nn.Linear(10, 10),
|
||||
nn.ReLU(),
|
||||
nn.Linear(10, 10),
|
||||
nn.ReLU(),
|
||||
nn.Linear(10, 10),
|
||||
nn.ReLU(),
|
||||
nn.Linear(10, 1),
|
||||
)
|
||||
|
||||
def forward(self, obs, acts):
|
||||
# obs.shape = (batch=514, 1 + others=150, 5)
|
||||
# acts.shape = (batch=514, 1)
|
||||
# val.shape = (batch=514,)
|
||||
|
||||
ego = obs[:, 0]
|
||||
rel = obs[:, 1:]
|
||||
nan = rel.isnan().any(-1, keepdim=True)
|
||||
rel = torch.where(nan, torch.zeros_like(rel), rel) # required because of https://github.com/pytorch/pytorch/issues/15506
|
||||
|
||||
d = (rel[:, :, :2] ** 2).sum(-1).sqrt()
|
||||
front = torch.stack((torch.cos(ego[:, 3]), torch.sin(ego[:, 3])), -1)
|
||||
left = torch.stack((-torch.sin(ego[:, 3]), torch.cos(ego[:, 3])), -1)
|
||||
df = (rel[:, :, :2] * front.unsqueeze(1)).sum(-1)
|
||||
dl = (rel[:, :, :2] * left.unsqueeze(1)).sum(-1)
|
||||
alpha = torch.atan2(dl, df)
|
||||
|
||||
rel[:, :, 0] = d
|
||||
rel[:, :, 1] = alpha
|
||||
|
||||
e = self._ego_encoder(ego)
|
||||
x = self._state_encoder(rel)
|
||||
x = torch.where(nan, torch.zeros_like(x), x)
|
||||
x = self._deepset(x)
|
||||
v = self._discriminator(torch.cat((e, x, acts), 1))
|
||||
|
||||
return v.squeeze(1)
|
||||
31
scratch/etienne/pillbox/learners/soft_q.py
Normal file
31
scratch/etienne/pillbox/learners/soft_q.py
Normal file
@@ -0,0 +1,31 @@
|
||||
from typing import Any, Dict, List, Optional, Type
|
||||
|
||||
import gym
|
||||
import torch as th
|
||||
from torch import nn
|
||||
|
||||
from stable_baselines3.common.policies import BasePolicy, register_policy
|
||||
from stable_baselines3.common.torch_layers import BaseFeaturesExtractor, FlattenExtractor, NatureCNN, create_mlp
|
||||
from stable_baselines3.dqn.policies import DQNPolicy, QNetwork
|
||||
|
||||
|
||||
class SoftQNetwork(QNetwork):
|
||||
def _predict(self, observation: th.Tensor, deterministic: bool = True) -> th.Tensor:
|
||||
q_values = self.forward(observation)
|
||||
probs = nn.functional.softmax(q_values * 10, dim=1)
|
||||
m = th.distributions.Categorical(probs)
|
||||
action = m.sample().reshape(-1)
|
||||
return action
|
||||
|
||||
|
||||
class SQLPolicy(DQNPolicy):
|
||||
def make_q_net(self) -> SoftQNetwork:
|
||||
# Make sure we always have separate networks for features extractors etc
|
||||
net_args = self._update_features_extractor(
|
||||
self.net_args, features_extractor=None)
|
||||
return SoftQNetwork(**net_args).to(self.device)
|
||||
|
||||
|
||||
SoftMlpPolicy = SQLPolicy
|
||||
|
||||
register_policy("SoftMlpPolicy", SoftMlpPolicy)
|
||||
61
scratch/etienne/pillbox/learners/sqil.py
Normal file
61
scratch/etienne/pillbox/learners/sqil.py
Normal file
@@ -0,0 +1,61 @@
|
||||
import warnings
|
||||
from abc import ABC, abstractmethod
|
||||
from typing import Dict, Generator, Optional, Union
|
||||
|
||||
import numpy as np
|
||||
import torch as th
|
||||
from gym import spaces
|
||||
|
||||
try:
|
||||
# Check memory used by replay buffer when possible
|
||||
import psutil
|
||||
except ImportError:
|
||||
psutil = None
|
||||
|
||||
from stable_baselines3.common.preprocessing import get_action_dim, get_obs_shape
|
||||
from stable_baselines3.common.type_aliases import ReplayBufferSamples, RolloutBufferSamples
|
||||
from stable_baselines3.common.vec_env import VecNormalize
|
||||
from stable_baselines3.common.buffers import ReplayBuffer
|
||||
|
||||
|
||||
class SQILReplayBuffer(ReplayBuffer):
|
||||
def __init__(
|
||||
self,
|
||||
buffer_size: int,
|
||||
observation_space: spaces.Space,
|
||||
action_space: spaces.Space,
|
||||
device: Union[th.device, str] = "cpu",
|
||||
n_envs: int = 1,
|
||||
optimize_memory_usage: bool = False,
|
||||
expert_data: dict = dict(),
|
||||
):
|
||||
super(SQILReplayBuffer, self).__init__(buffer_size, observation_space, action_space, device, n_envs=n_envs, optimize_memory_usage=optimize_memory_usage)
|
||||
|
||||
self.expert_states = expert_data['obs']
|
||||
self.expert_actions = expert_data['acts']
|
||||
self.expert_next_states = expert_data['next_obs']
|
||||
self.expert_dones = expert_data['dones']
|
||||
|
||||
def _get_samples(self, batch_inds: np.ndarray, env: Optional[VecNormalize] = None) -> ReplayBufferSamples:
|
||||
num_samples = len(batch_inds)
|
||||
num_expert_samples = int(num_samples / 2)
|
||||
batch_inds = batch_inds[:num_expert_samples]
|
||||
expert_inds = np.random.randint(0, len(self.expert_states), size=num_expert_samples)
|
||||
# Balanced sampling
|
||||
if self.optimize_memory_usage:
|
||||
next_obs = self._normalize_obs(self.observations[(batch_inds + 1) % self.buffer_size, 0, :], env)
|
||||
else:
|
||||
next_obs = self._normalize_obs(self.next_observations[batch_inds, 0, :], env)
|
||||
next_obs = np.concatenate((next_obs, self._normalize_obs(self.expert_next_states[expert_inds], env)), axis=0)
|
||||
obs = self._normalize_obs(self.observations[batch_inds, 0, :], env)
|
||||
obs = np.concatenate((obs, self._normalize_obs(self.expert_states[expert_inds], env)), axis=0)
|
||||
actions = self.actions[batch_inds, 0, :]
|
||||
actions = np.concatenate((actions, self.expert_actions[expert_inds].reshape(num_expert_samples, -1)), axis=0)
|
||||
dones = self.dones[batch_inds]
|
||||
dones = np.concatenate((dones, self.expert_dones[expert_inds].reshape(num_expert_samples, -1)), axis=0)
|
||||
# SQIL Rewards
|
||||
rewards = self.rewards[batch_inds] * 0.
|
||||
rewards = np.concatenate((rewards, np.ones_like(rewards)), axis=0)
|
||||
|
||||
data = (obs, actions, next_obs, dones, rewards)
|
||||
return ReplayBufferSamples(*tuple(map(self.to_torch, data)))
|
||||
248
scratch/etienne/pillbox/learners/train.py
Normal file
248
scratch/etienne/pillbox/learners/train.py
Normal file
@@ -0,0 +1,248 @@
|
||||
from imitation.algorithms import adversarial, bc
|
||||
from imitation.util import logger, util
|
||||
from stable_baselines3 import PPO, DQN, SAC
|
||||
from soft_q import SQLPolicy
|
||||
from sqil import SQILReplayBuffer
|
||||
from stable_baselines3.common import policies
|
||||
from stable_baselines3.common.evaluation import evaluate_policy
|
||||
from imitation.rewards import discrim_nets
|
||||
import numpy as np
|
||||
import argparse
|
||||
from utils import make_sa_dataloader, make_sads_dataloader, make_sa_dataset, linear_schedule
|
||||
from stable_baselines3.common.vec_env import DummyVecEnv, VecNormalize
|
||||
from adril import AdRILWrapper, AdRILReplayBuffer
|
||||
import os
|
||||
from gym.spaces import Discrete
|
||||
import gym
|
||||
from advil import advil_training
|
||||
from stable_baselines3.common.running_mean_std import RunningMeanStd
|
||||
|
||||
from advil import AdVILPolicy, AdVILDiscriminator
|
||||
|
||||
def train_bc(env, n=0):
|
||||
venv = util.make_vec_env(env, n_envs=8)
|
||||
if isinstance(venv.action_space, Discrete):
|
||||
w = 64
|
||||
else:
|
||||
w = 256
|
||||
for i in range(n):
|
||||
mean_rewards = []
|
||||
std_rewards = []
|
||||
for num_trajs in range(0, 26, 5):
|
||||
if num_trajs == 0:
|
||||
expert_data = make_sa_dataloader(env, normalize=False)
|
||||
else:
|
||||
expert_data = make_sa_dataloader(env, max_trajs=num_trajs, normalize=False)
|
||||
bc_trainer = bc.BC(venv.observation_space, venv.action_space, expert_data=expert_data,
|
||||
policy_class=policies.ActorCriticPolicy,
|
||||
ent_weight=0., l2_weight=0., policy_kwargs=dict(net_arch=[w, w]))
|
||||
if num_trajs > 0:
|
||||
bc_trainer.train(n_batches=int(5e5))
|
||||
|
||||
def get_policy(*args, **kwargs):
|
||||
return bc_trainer.policy
|
||||
model = PPO(get_policy, env, verbose=1)
|
||||
model.save(os.path.join("learners", env,
|
||||
"bc_{0}_{1}".format(i, num_trajs)))
|
||||
mean_reward, std_reward = evaluate_policy(
|
||||
model, model.get_env(), n_eval_episodes=10)
|
||||
mean_rewards.append(mean_reward)
|
||||
std_rewards.append(std_reward)
|
||||
print("{0} Trajs: {1}".format(num_trajs, mean_reward))
|
||||
np.savez(os.path.join("learners", env, "bc_rewards_{0}".format(
|
||||
i)), means=mean_rewards, stds=std_rewards)
|
||||
|
||||
|
||||
def train_gail(env, n=0):
|
||||
venv = util.make_vec_env(env, n_envs=8)
|
||||
if isinstance(venv.action_space, Discrete):
|
||||
w = 64
|
||||
else:
|
||||
w = 256
|
||||
expert_data = make_sads_dataloader(env, max_trajs=5)
|
||||
logger.configure(os.path.join("learners", "GAIL"))
|
||||
|
||||
for i in range(n):
|
||||
discrim_net = discrim_nets.ActObsMLP(
|
||||
action_space=venv.action_space,
|
||||
observation_space=venv.observation_space,
|
||||
hid_sizes=(w, w),
|
||||
)
|
||||
gail_trainer = adversarial.GAIL(venv, expert_data=expert_data, expert_batch_size=32,
|
||||
gen_algo=PPO("MlpPolicy", venv, verbose=1, n_steps=1024,
|
||||
policy_kwargs=dict(net_arch=[w, w])),
|
||||
discrim_kwargs={'discrim_net': discrim_net})
|
||||
mean_rewards = []
|
||||
std_rewards = []
|
||||
for train_steps in range(20):
|
||||
if train_steps > 0:
|
||||
if 'Bullet' in env:
|
||||
gail_trainer.train(total_timesteps=25000)
|
||||
else:
|
||||
gail_trainer.train(total_timesteps=16384)
|
||||
|
||||
def get_policy(*args, **kwargs):
|
||||
return gail_trainer.gen_algo.policy
|
||||
model = PPO(get_policy, env, verbose=1)
|
||||
mean_reward, std_reward = evaluate_policy(
|
||||
model, model.env, n_eval_episodes=10)
|
||||
mean_rewards.append(mean_reward)
|
||||
std_rewards.append(std_reward)
|
||||
print("{0} Steps: {1}".format(train_steps, mean_reward))
|
||||
np.savez(os.path.join("learners", env, "gail_rewards_{0}".format(i)),
|
||||
means=mean_rewards, stds=std_rewards)
|
||||
|
||||
|
||||
def train_sqil(env, n=0):
|
||||
venv = gym.make(env)
|
||||
expert_data = make_sa_dataset(env, max_trajs=5)
|
||||
|
||||
for i in range(n):
|
||||
if isinstance(venv.action_space, Discrete):
|
||||
model = DQN(SQLPolicy, venv, verbose=1, policy_kwargs=dict(net_arch=[64, 64]), learning_starts=1)
|
||||
else:
|
||||
model = SAC('MlpPolicy', venv, verbose=1, policy_kwargs=dict(net_arch=[256, 256]), ent_coef='auto',
|
||||
learning_rate=linear_schedule(7.3e-4), train_freq=64, gradient_steps=64, gamma=0.98, tau=0.02)
|
||||
|
||||
model.replay_buffer = SQILReplayBuffer(model.buffer_size, model.observation_space,
|
||||
model.action_space, model.device, 1,
|
||||
model.optimize_memory_usage, expert_data=expert_data)
|
||||
mean_rewards = []
|
||||
std_rewards = []
|
||||
for train_steps in range(20):
|
||||
if train_steps > 0:
|
||||
if 'Bullet' in env:
|
||||
model.learn(total_timesteps=25000, log_interval=1)
|
||||
else:
|
||||
model.learn(total_timesteps=16384, log_interval=1)
|
||||
mean_reward, std_reward = evaluate_policy(
|
||||
model, model.env, n_eval_episodes=10)
|
||||
mean_rewards.append(mean_reward)
|
||||
std_rewards.append(std_reward)
|
||||
print("{0} Steps: {1}".format(train_steps, mean_reward))
|
||||
np.savez(os.path.join("learners", env, "sqil_rewards_{0}".format(i)),
|
||||
means=mean_rewards, stds=std_rewards)
|
||||
|
||||
|
||||
def train_adril(env, n=0, balanced=False):
|
||||
num_trajs = 20
|
||||
expert_data = make_sa_dataset(env, max_trajs=num_trajs)
|
||||
n_expert = len(expert_data["obs"])
|
||||
expert_sa = np.concatenate((expert_data["obs"], np.reshape(expert_data["acts"], (n_expert, -1))), axis=1)
|
||||
|
||||
for i in range(0, n):
|
||||
venv = AdRILWrapper(gym.make(env))
|
||||
mean_rewards = []
|
||||
std_rewards = []
|
||||
# Create model
|
||||
if isinstance(venv.action_space, Discrete):
|
||||
model = DQN(SQLPolicy, venv, verbose=1, policy_kwargs=dict(net_arch=[64, 64]), learning_starts=1)
|
||||
else:
|
||||
model = SAC('MlpPolicy', venv, verbose=1, policy_kwargs=dict(net_arch=[256, 256]), ent_coef='auto',
|
||||
learning_rate=linear_schedule(7.3e-4), train_freq=64, gradient_steps=64, gamma=0.98, tau=0.02)
|
||||
model.replay_buffer = AdRILReplayBuffer(model.buffer_size, model.observation_space,
|
||||
model.action_space, model.device, 1,
|
||||
model.optimize_memory_usage, expert_data=expert_data, N_expert=num_trajs,
|
||||
balanced=balanced)
|
||||
if not balanced:
|
||||
for j in range(len(expert_sa)):
|
||||
obs = expert_data["obs"][j]
|
||||
act = expert_data["acts"][j]
|
||||
next_obs = expert_data["next_obs"][j]
|
||||
done = expert_data["dones"][j]
|
||||
model.replay_buffer.add(obs, next_obs, act, -1, done)
|
||||
for train_steps in range(400):
|
||||
# Train policy
|
||||
if train_steps > 0:
|
||||
if 'Bullet' in env:
|
||||
model.learn(total_timesteps=1250, log_interval=1000)
|
||||
else:
|
||||
model.learn(total_timesteps=25000, log_interval=1000)
|
||||
if train_steps % 1 == 0: # written to support more complex update schemes
|
||||
model.replay_buffer.set_iter(train_steps)
|
||||
model.replay_buffer.set_n_learner(venv.num_trajs)
|
||||
|
||||
# Evaluate policy
|
||||
if train_steps % 20 == 0:
|
||||
model.set_env(gym.make(env))
|
||||
mean_reward, std_reward = evaluate_policy(
|
||||
model, model.env, n_eval_episodes=10)
|
||||
mean_rewards.append(mean_reward)
|
||||
std_rewards.append(std_reward)
|
||||
print("{0} Steps: {1}".format(int(train_steps * 1250), mean_reward))
|
||||
np.savez(os.path.join("learners", env, "adril_rewards_{0}".format(i)),
|
||||
means=mean_rewards, stds=std_rewards)
|
||||
# Update env
|
||||
if train_steps > 0:
|
||||
if train_steps % 1 == 0:
|
||||
venv.set_iter(train_steps + 1)
|
||||
model.set_env(venv)
|
||||
|
||||
|
||||
def train_advil(env, policy_class=AdVILPolicy, discriminator_class=AdVILDiscriminator,
|
||||
iters=int(1e5), lr_pi=8e-6, lr_f=8e-4):
|
||||
venv = gym.make(env)
|
||||
expert_data = make_sa_dataloader(
|
||||
env,
|
||||
normalize=False,
|
||||
batch_size=1024,
|
||||
)
|
||||
pi = advil_training(
|
||||
expert_data,
|
||||
venv,
|
||||
iters=iters,
|
||||
policy_class=policy_class,
|
||||
discriminator_class=discriminator_class,
|
||||
lr_pi=lr_pi,
|
||||
lr_f=lr_f,
|
||||
)
|
||||
return pi
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = argparse.ArgumentParser(description='Train expert policies.')
|
||||
parser.add_argument(
|
||||
'-a', '--algo', choices=['bc', 'gail', 'sqil', 'adril', 'advil', 'all'], required=True)
|
||||
parser.add_argument('-e', '--env', choices=['cartpole', 'lunarlander', 'acrobot', 'pendulum', 'halfcheetah', 'walker', 'hopper', 'ant'],
|
||||
required=True)
|
||||
parser.add_argument('-n', '--num_runs', required=False)
|
||||
args = parser.parse_args()
|
||||
if args.env == "cartpole":
|
||||
envname = 'CartPole-v1'
|
||||
elif args.env == "lunarlander":
|
||||
envname = 'LunarLander-v2'
|
||||
elif args.env == "acrobot":
|
||||
envname = 'Acrobot-v1'
|
||||
elif args.env == "pendulum":
|
||||
envname = 'Pendulum-v0'
|
||||
elif args.env == "halfcheetah":
|
||||
envname = 'HalfCheetahBulletEnv-v0'
|
||||
elif args.env == "walker":
|
||||
envname = 'Walker2DBulletEnv-v0'
|
||||
elif args.env == "hopper":
|
||||
envname = 'HopperBulletEnv-v0'
|
||||
elif args.env == "ant":
|
||||
envname = 'AntBulletEnv-v0'
|
||||
else:
|
||||
print("ERROR: unsupported env.")
|
||||
if args.num_runs is not None and args.num_runs.isdigit():
|
||||
num_runs = int(args.num_runs)
|
||||
else:
|
||||
num_runs = 1
|
||||
if args.algo == 'bc':
|
||||
train_bc(envname, num_runs)
|
||||
elif args.algo == 'gail':
|
||||
train_gail(envname, num_runs)
|
||||
elif args.algo == 'sqil':
|
||||
train_sqil(envname, num_runs)
|
||||
elif args.algo == 'adril':
|
||||
train_adril(envname, num_runs)
|
||||
elif args.algo == 'advil':
|
||||
train_advil(envname, num_runs)
|
||||
elif args.algo == 'all':
|
||||
train_bc(envname, num_runs)
|
||||
train_gail(envname, num_runs)
|
||||
train_sqil(envname, num_runs)
|
||||
train_adril(envname, num_runs)
|
||||
train_advil(envname, num_runs)
|
||||
else:
|
||||
print("ERROR: unsupported algorithm")
|
||||
129
scratch/etienne/pillbox/learners/utils.py
Normal file
129
scratch/etienne/pillbox/learners/utils.py
Normal file
@@ -0,0 +1,129 @@
|
||||
import numpy as np
|
||||
import torch
|
||||
from torch.utils.data import Dataset, DataLoader
|
||||
from itertools import chain
|
||||
from typing import Callable, Union, Type, Optional, Dict, Any
|
||||
|
||||
# From https://github.com/DLR-RM/rl-baselines3-zoo/blob/8ea4f4a87afa548832ca17e575b351ec5928c1b0/utils/utils.py
|
||||
def linear_schedule(initial_value: Union[float, str]) -> Callable[[float], float]:
|
||||
"""
|
||||
Linear learning rate schedule.
|
||||
:param initial_value: (float or str)
|
||||
:return: (function)
|
||||
"""
|
||||
if isinstance(initial_value, str):
|
||||
initial_value = float(initial_value)
|
||||
|
||||
def func(progress_remaining: float) -> float:
|
||||
"""
|
||||
Progress will decrease from 1 (beginning) to 0
|
||||
:param progress_remaining: (float)
|
||||
:return: (float)
|
||||
"""
|
||||
return progress_remaining * initial_value
|
||||
|
||||
return func
|
||||
|
||||
class SADataset(torch.utils.data.Dataset):
|
||||
def __init__(self, obs, acts, normalize):
|
||||
if normalize:
|
||||
obs = np.array(obs)
|
||||
self.mean = obs.mean(axis=0)
|
||||
self.std = obs.std(axis=0) + 1e-3
|
||||
obs = (obs - self.mean) / (self.std)
|
||||
self.is_normalized = True
|
||||
else:
|
||||
self.is_normalized = False
|
||||
self.obs = torch.tensor(obs)
|
||||
self.acts = torch.tensor(acts)
|
||||
|
||||
def __len__(self):
|
||||
return len(self.obs)
|
||||
|
||||
def __getitem__(self, idx):
|
||||
if torch.is_tensor(idx):
|
||||
idx = idx.tolist()
|
||||
obs = self.obs[idx]
|
||||
acts = self.acts[idx]
|
||||
sample = {'obs': obs, 'acts': acts}
|
||||
return sample
|
||||
|
||||
def make_sa_dataloader(envname, max_trajs=None, normalize=False, batch_size=32):
|
||||
demos = np.load(
|
||||
"../experts/{0}/demos.npz".format(envname), allow_pickle=True)
|
||||
num_trajs = demos["num_trajs"]
|
||||
if max_trajs is None:
|
||||
max_trajs = num_trajs
|
||||
obs = []
|
||||
acts = []
|
||||
for traj in range(min(max_trajs, num_trajs)):
|
||||
obs.extend(demos[str(traj)].item()['states'])
|
||||
acts.extend(demos[str(traj)].item()['actions'])
|
||||
dataset = SADataset(obs, acts, normalize)
|
||||
dataloader = DataLoader(dataset, batch_size=batch_size,
|
||||
shuffle=True, num_workers=0)
|
||||
return dataloader
|
||||
|
||||
class SADSDataset(torch.utils.data.Dataset):
|
||||
def __init__(self, obs, acts, next_obs, traj_lens):
|
||||
self.obs = torch.tensor(obs)
|
||||
self.acts = torch.tensor(acts)
|
||||
self.next_obs = torch.tensor(next_obs)
|
||||
dones = [[False for _ in range(l - 2)] + [True] for l in traj_lens]
|
||||
self.dones = torch.tensor(list(chain.from_iterable(dones)))
|
||||
|
||||
def __len__(self):
|
||||
return len(self.obs)
|
||||
|
||||
def __getitem__(self, idx):
|
||||
if torch.is_tensor(idx):
|
||||
idx = idx.tolist()
|
||||
obs = self.obs[idx]
|
||||
acts = self.acts[idx]
|
||||
next_obs = self.next_obs[idx]
|
||||
dones = self.dones[idx]
|
||||
sample = {'obs': obs, 'acts': acts,
|
||||
'next_obs': next_obs, 'dones': dones}
|
||||
return sample
|
||||
|
||||
def make_sads_dataloader(envname, max_trajs=None):
|
||||
demos = np.load(
|
||||
"./experts/{0}/demos.npz".format(envname), allow_pickle=True)
|
||||
num_trajs = demos["num_trajs"]
|
||||
if max_trajs is None:
|
||||
max_trajs = num_trajs
|
||||
obs = []
|
||||
next_obs = []
|
||||
acts = []
|
||||
lens = []
|
||||
for traj in range(min(max_trajs, num_trajs)):
|
||||
obs.extend(demos[str(traj)].item()['states'][:-1])
|
||||
next_obs.extend(demos[str(traj)].item()['states'][1:])
|
||||
acts.extend(demos[str(traj)].item()['actions'][:-1])
|
||||
lens.append(len(demos[str(traj)].item()['states']))
|
||||
dataset = SADSDataset(obs, acts, next_obs, lens)
|
||||
dataloader = DataLoader(dataset, batch_size=32,
|
||||
shuffle=False, num_workers=0, drop_last=True)
|
||||
return dataloader
|
||||
|
||||
def make_sa_dataset(envname, max_trajs=None):
|
||||
demos = np.load("../pillbox/experts/{0}/demos.npz".format(envname), allow_pickle=True)
|
||||
num_trajs = demos["num_trajs"]
|
||||
if max_trajs is None:
|
||||
max_trajs = num_trajs
|
||||
expert_states = []
|
||||
expert_actions = []
|
||||
expert_next_states = []
|
||||
expert_dones = []
|
||||
for traj in range(min(max_trajs, num_trajs)):
|
||||
expert_states.extend(demos[str(traj)].item()['states'][:-1])
|
||||
expert_next_states.extend(demos[str(traj)].item()['states'][1:])
|
||||
expert_actions.extend(demos[str(traj)].item()['actions'][:-1])
|
||||
l = len(demos[str(traj)].item()['states'])
|
||||
expert_dones.extend([False for _ in range(l - 2)] + [True])
|
||||
expert_data = dict()
|
||||
expert_data['obs'] = np.array(expert_states)
|
||||
expert_data['acts'] = np.array(expert_actions)
|
||||
expert_data['next_obs'] = np.array(expert_next_states)
|
||||
expert_data['dones'] = np.array(expert_dones)
|
||||
return expert_data
|
||||
9
scratch/etienne/pillbox/requirements.txt
Normal file
9
scratch/etienne/pillbox/requirements.txt
Normal file
@@ -0,0 +1,9 @@
|
||||
gym
|
||||
numpy
|
||||
psutil
|
||||
scikit_learn
|
||||
scipy
|
||||
stable_baselines3
|
||||
torch
|
||||
tqdm
|
||||
imitation
|
||||
@@ -123,7 +123,7 @@ def train(config, policy, train_dataset, cv_dataset, filestr, **kwargs):
|
||||
cv_loader = DataLoader(cv_dataset, batch_size=cv_batch_size, shuffle=True)
|
||||
|
||||
# change policy dtype
|
||||
policy.policy = policy.policy.type(train_dataset[0]['state'].dtype)
|
||||
policy.policy = policy.policy.type(train_dataset[0]['state']['ego_state'].dtype)
|
||||
|
||||
# generate loss function, optimizer
|
||||
cv_loss_fn = nn.MSELoss(reduction='sum')
|
||||
@@ -150,7 +150,7 @@ def train(config, policy, train_dataset, cv_dataset, filestr, **kwargs):
|
||||
for (batch_idx, batch) in enumerate(training_loader):
|
||||
|
||||
# sample mini-batch and run through policy
|
||||
pred_action = policy(batch)
|
||||
pred_action = policy(batch['state'])
|
||||
loss = loss_fn(pred_action, batch['action'])
|
||||
|
||||
# compute loss and step optimizer
|
||||
@@ -169,7 +169,7 @@ def train(config, policy, train_dataset, cv_dataset, filestr, **kwargs):
|
||||
with torch.no_grad():
|
||||
cv_loss = 0.
|
||||
for (batch_idx, batch) in enumerate(cv_loader):
|
||||
pred_action = policy(batch)
|
||||
pred_action = policy(batch['state'])
|
||||
loss = cv_loss_fn(pred_action, batch['action'])
|
||||
cv_loss += loss.item() / len(cv_dataset)
|
||||
|
||||
|
||||
@@ -25,13 +25,14 @@ class InteractionDatasetSingleAgent(Dataset):
|
||||
self.loc = loc
|
||||
self.tracks = tracks
|
||||
self.dtype = dtype
|
||||
self.keys = ['ego_state', 'relative_state', 'path', 'action', 'next_ego_state', 'next_relative_state', 'next_path']
|
||||
self._load_dataset()
|
||||
|
||||
def _load_dataset(self):
|
||||
"""
|
||||
Load the full datasets ahead of time
|
||||
"""
|
||||
self.raw_data = {'state':[], 'relative_state':[], 'action':[], 'path_x':[], 'path_y':[]}
|
||||
self.raw_data = {key:[] for key in self.keys}
|
||||
max_nv = 0
|
||||
for track in self.tracks:
|
||||
try:
|
||||
@@ -41,33 +42,26 @@ class InteractionDatasetSingleAgent(Dataset):
|
||||
print('Failed to load location {} track {}'.format(self.loc,track))
|
||||
continue
|
||||
max_nv = max(max_nv, data['relative_state'].shape[1])
|
||||
self.raw_data['state'].append(data['state'])
|
||||
self.raw_data['relative_state'].append(data['relative_state'])
|
||||
self.raw_data['action'].append(data['action'])
|
||||
self.raw_data['path_x'].append(data['path_x'])
|
||||
self.raw_data['path_y'].append(data['path_y'])
|
||||
|
||||
# cat lists
|
||||
self.raw_data['state'] = torch.cat(self.raw_data['state']).type(self.dtype)
|
||||
self.raw_data['action'] = torch.cat(self.raw_data['action']).type(self.dtype)
|
||||
self.raw_data['path_x'] = torch.cat(self.raw_data['path_x']).type(self.dtype)
|
||||
self.raw_data['path_y'] = torch.cat(self.raw_data['path_y']).type(self.dtype)
|
||||
for key in self.keys:
|
||||
self.raw_data[key].append(data[key])
|
||||
|
||||
# pad second dimension of relative state
|
||||
for i in range(len(self.raw_data['relative_state'])):
|
||||
nv1, nv2, d = self.raw_data['relative_state'][i].shape
|
||||
pad = torch.zeros(nv1, max_nv-nv2, d, dtype=self.dtype) * np.nan
|
||||
self.raw_data['relative_state'][i] = torch.cat((self.raw_data['relative_state'][i], pad), dim=1)
|
||||
self.raw_data['relative_state'] = torch.cat(self.raw_data['relative_state']).type(self.dtype)
|
||||
self.raw_data['next_relative_state'][i] = torch.cat((self.raw_data['next_relative_state'][i], pad), dim=1)
|
||||
|
||||
# cat lists
|
||||
for key in self.keys:
|
||||
self.raw_data[key] = torch.cat(self.raw_data[key]).type(self.dtype)
|
||||
|
||||
# mandate equal length
|
||||
assert len(self.raw_data['state']) == len(self.raw_data['relative_state']) \
|
||||
== len(self.raw_data['action']) \
|
||||
== len(self.raw_data['path_x']) \
|
||||
== len(self.raw_data['path_y']), 'dataset lengths unequal'
|
||||
lengths = [len(self.raw_data[key]) for key in self.keys]
|
||||
assert min(lengths) == max(lengths), 'dataset lengths unequal'
|
||||
|
||||
def __len__(self):
|
||||
return len(self.raw_data['state'])
|
||||
return len(self.raw_data['ego_state'])
|
||||
|
||||
def __getitem__(self, idx):
|
||||
"""
|
||||
@@ -76,12 +70,27 @@ class InteractionDatasetSingleAgent(Dataset):
|
||||
idx: index or indices of B samples
|
||||
Returns:
|
||||
sample (dict): sample dictionary with the following entries:
|
||||
state (torch.tensor): (B, 5) raw state
|
||||
state (dict): state dictionary with the following entries:
|
||||
ego_state (torch.tensor): (B, 5) raw state
|
||||
relative_state (torch.tensor): (B, max_nv, d) relative state (padded with nans)
|
||||
path_x (torch.tensor): (B, P) tensor of P future path x positions
|
||||
path_y (torch.tensor): (B, P) tensor of P future path y positions
|
||||
path (torch.tensor): (B, P, 2) tensor of P future path x and y positions
|
||||
action (torch.tensor): (B, 1) actions taken from each state
|
||||
next_stat (dict): next state dictionary with the following entries:
|
||||
ego_state (torch.tensor): (B, 5) raw next state
|
||||
relative_state (torch.tensor): (B, max_nv, d) next relative state (padded with nans)
|
||||
path (torch.tensor): (B, P, 2) tensor of P future next path x and y positions
|
||||
"""
|
||||
keys = ['state', 'relative_state', 'path_x', 'path_y', 'action']
|
||||
sample = {key:self.raw_data[key][idx] for key in keys}
|
||||
#sample = {key:self.raw_data[key][idx] for key in self.keys}
|
||||
sample = {
|
||||
'state':{
|
||||
'ego_state':self.raw_data['ego_state'][idx],
|
||||
'relative_state':self.raw_data['relative_state'][idx],
|
||||
'path':self.raw_data['path'][idx]
|
||||
},
|
||||
'action':self.raw_data['action'][idx],
|
||||
'next_state':{
|
||||
'ego_state':self.raw_data['next_ego_state'][idx],
|
||||
'relative_state':self.raw_data['next_relative_state'][idx],
|
||||
'path':self.raw_data['next_path'][idx]},
|
||||
}
|
||||
return sample
|
||||
@@ -31,8 +31,8 @@ def generate_expert_data(path: str='expert_data', loc: int = 0, track:int = 0,
|
||||
os.makedirs(path)
|
||||
filestr = opj(path,intersim.LOCATIONS[loc]+'_track%03i'%(track))
|
||||
|
||||
svt, svt_path = get_svt(base='InteractionSimulator', loc=loc, track=track)
|
||||
osm = get_map_path(base='InteractionSimulator', loc=loc)
|
||||
svt, svt_path = get_svt(loc=loc, track=track) #base='InteractionSimulator'
|
||||
osm = get_map_path(loc=loc)
|
||||
print('SVT path: {}'.format(svt_path))
|
||||
print('Map path: {}'.format(osm))
|
||||
states, actions = SVT_to_stateactions(svt)
|
||||
@@ -91,33 +91,42 @@ def process_expert_observations(obs, actions, filestr, remove_outliers=True, dty
|
||||
remove_outliers (bool): whether to remove datapoints with acceleration above or below 5 m/s/s
|
||||
dtype (torch.Type): type to convert data to
|
||||
"""
|
||||
keys = ['state', 'action', 'relative_state', 'path_x', 'path_y']
|
||||
keys = ['ego_state', 'relative_state', 'path', 'action', 'next_ego_state', 'next_relative_state', 'next_path']
|
||||
data = {key:[] for key in keys}
|
||||
assert len(obs) == len(actions), 'non-matching action and observation lengths'
|
||||
T = len(obs)
|
||||
max_nv = 0
|
||||
for t in range(T):
|
||||
nni = ~torch.isnan(obs[t]['state'][:,0])
|
||||
for t in range(T-1):
|
||||
nni = ~torch.isnan(obs[t]['state'][:,0]) & ~torch.isnan(obs[t+1]['state'][:,0])
|
||||
max_nv = max(max_nv,nni.count_nonzero())
|
||||
data['state'].append(obs[t]['state'][nni])
|
||||
|
||||
# state
|
||||
data['ego_state'].append(obs[t]['state'][nni])
|
||||
data['relative_state'].append(obs[t]['relative_state'].index_select(0,
|
||||
nni.nonzero()[:,0]).index_select(1, nni.nonzero()[:,0]))
|
||||
data['action'].append(actions[t][nni])
|
||||
data['path_x'].append(obs[t]['paths'][0][nni])
|
||||
data['path_y'].append(obs[t]['paths'][1][nni])
|
||||
data['path'].append(torch.stack((obs[t]['paths'][0][nni], obs[t]['paths'][1][nni]), dim=-1))
|
||||
|
||||
# action
|
||||
data['action'].append(actions[t][nni])
|
||||
|
||||
# next state
|
||||
data['next_ego_state'].append(obs[t+1]['state'][nni])
|
||||
data['next_relative_state'].append(obs[t+1]['relative_state'].index_select(0,
|
||||
nni.nonzero()[:,0]).index_select(1, nni.nonzero()[:,0]))
|
||||
data['next_path'].append(torch.stack((obs[t+1]['paths'][0][nni], obs[t+1]['paths'][1][nni]), dim=-1))
|
||||
|
||||
|
||||
# cat lists
|
||||
data['state'] = torch.cat(data['state']).type(dtype)
|
||||
data['action'] = torch.cat(data['action']).type(dtype)
|
||||
data['path_x'] = torch.cat(data['path_x']).type(dtype)
|
||||
data['path_y'] = torch.cat(data['path_y']).type(dtype)
|
||||
|
||||
# pad second dimension of relative state
|
||||
for i in range(len(data['relative_state'])):
|
||||
nv1, nv2, d = data['relative_state'][i].shape
|
||||
pad = torch.zeros(nv1, max_nv-nv2, d, dtype=dtype) * np.nan
|
||||
data['relative_state'][i] = torch.cat((data['relative_state'][i], pad), dim=1)
|
||||
data['relative_state'] = torch.cat(data['relative_state']).type(dtype)
|
||||
data['next_relative_state'][i] = torch.cat((data['next_relative_state'][i], pad), dim=1)
|
||||
|
||||
# cat lists
|
||||
for key in keys:
|
||||
data[key] = torch.cat(data[key]).type(dtype)
|
||||
|
||||
if remove_outliers:
|
||||
non_outlier_indices = torch.nonzero(torch.abs(data['action'][:,0]) < 5)
|
||||
@@ -125,12 +134,11 @@ def process_expert_observations(obs, actions, filestr, remove_outliers=True, dty
|
||||
data[key] = data[key][non_outlier_indices[:,0]]
|
||||
|
||||
# mandate equal length
|
||||
assert len(data['state']) == len(data['relative_state']) \
|
||||
== len(data['action']) == len(data['path_x']) \
|
||||
== len(data['path_y']), 'dataset lengths unequal'
|
||||
lengths = [len(data[key]) for key in keys]
|
||||
assert min(lengths) == max(lengths), 'dataset lengths unequal'
|
||||
|
||||
# save out data
|
||||
for key in data.keys():
|
||||
for key in keys:
|
||||
torch.save(data[key], filestr+'_'+key+'.pt')
|
||||
|
||||
def load_expert_data(path='expert_data', loc: int = 0, track:int = 0):
|
||||
@@ -146,7 +154,8 @@ def load_expert_data(path='expert_data', loc: int = 0, track:int = 0):
|
||||
# load observations and actions
|
||||
filestr = opj(path, intersim.LOCATIONS[loc]+'_track%03i'%(track))
|
||||
data = {}
|
||||
for key in ['state','action','relative_state','path_x','path_y']:
|
||||
keys = ['ego_state', 'relative_state', 'path', 'action', 'next_ego_state', 'next_relative_state', 'next_path']
|
||||
for key in keys:
|
||||
data[key] = torch.load(filestr+'_'+key+'.pt')
|
||||
return data
|
||||
|
||||
|
||||
@@ -47,11 +47,11 @@ def main(config, method='bc', train=False, test=False, loc=0, datadir='./expert_
|
||||
if train:
|
||||
|
||||
# make policy, train and test datasets, and send to
|
||||
train_dataset = InteractionDatasetSingleAgent(output_dir=datadir, loc=loc, tracks=[0,1,2])
|
||||
train_dataset = InteractionDatasetSingleAgent(output_dir=datadir, loc=loc, tracks=kwargs['train_tracks'])
|
||||
# generate transform from train_dataset
|
||||
transforms = generate_transforms(train_dataset)
|
||||
policy = policy_class(config, transforms)
|
||||
cv_dataset = InteractionDatasetSingleAgent(output_dir=datadir, loc=loc, tracks=[3])
|
||||
cv_dataset = InteractionDatasetSingleAgent(output_dir=datadir, loc=loc, tracks=kwargs['cv_tracks'])
|
||||
train_fn(config, policy, train_dataset, cv_dataset, filestr, **kwargs)
|
||||
|
||||
if test:
|
||||
@@ -61,11 +61,10 @@ def main(config, method='bc', train=False, test=False, loc=0, datadir='./expert_
|
||||
policy.eval()
|
||||
|
||||
# simulate policy
|
||||
track = 4
|
||||
simulate_policy(policy, loc=loc, track=track, filestr=filestr, nframes=kwargs['nframes'], graph=kwargs['graph'])
|
||||
simulate_policy(policy, loc=loc, track=kwargs['test_tracks'][0], filestr=filestr, nframes=kwargs['nframes'], graph=kwargs['graph'])
|
||||
|
||||
# run test metrics
|
||||
test_dataset = InteractionDatasetSingleAgent(output_dir=datadir, loc=loc, tracks=[track])
|
||||
test_dataset = InteractionDatasetSingleAgent(output_dir=datadir, loc=loc, tracks=kwargs['test_tracks'])
|
||||
writer = SummaryWriter(filestr)
|
||||
info = metrics(filestr, test_dataset, policy)
|
||||
for k, m in info.items():
|
||||
|
||||
@@ -37,7 +37,7 @@ def metrics(filestr: str, test_dataset, policy):
|
||||
info['average_velocity'] = avg_v
|
||||
|
||||
# convert policy dtype between float32 and float64
|
||||
policy.policy = policy.policy.type(test_dataset[0]['state'].dtype)
|
||||
policy.policy = policy.policy.type(test_dataset[0]['state']['ego_state'].dtype)
|
||||
|
||||
# generate actions in test dataset
|
||||
true_actions, pred_actions = [], []
|
||||
@@ -45,9 +45,9 @@ def metrics(filestr: str, test_dataset, policy):
|
||||
test_loader = DataLoader(test_dataset, batch_size=1024)
|
||||
with torch.no_grad():
|
||||
for (batch_idx, batch) in enumerate(test_loader):
|
||||
pred_actions.append(policy(batch))
|
||||
pred_actions.append(policy(batch['state']))
|
||||
true_actions.append(batch['action'])
|
||||
true_velocities.append(batch['state'][:,2])
|
||||
true_velocities.append(batch['state']['ego_state'][:,2])
|
||||
|
||||
true_actions, pred_actions = torch.cat(true_actions,dim=0), torch.cat(pred_actions, dim=0)
|
||||
visualize_distribution(true_actions[:,0], pred_actions[:,0], filestr+'_action_viz')
|
||||
|
||||
@@ -31,17 +31,14 @@ class IntersimStateNet(nn.Module):
|
||||
sample (dict): sample dictionary with the following entries:
|
||||
state (torch.tensor): (B, 5) raw state
|
||||
relative_state (torch.tensor): (B, max_nv, d) relative state (padded with nans)
|
||||
path_x (torch.tensor): (B, P) tensor of P future path x positions
|
||||
path_y (torch.tensor): (B, P) tensor of P future path y positions
|
||||
path (torch.tensor): (B, P, 2) tensor of P future path x and y positions
|
||||
action (torch.tensor): (B, 1) actions taken from each state
|
||||
Returns:
|
||||
x (torch.tensor): (head_output_dim,) output of common head network
|
||||
"""
|
||||
ego = self.ego_net(sample["ego_state"])
|
||||
relative = self.deepsets_net(sample["relative_state"])
|
||||
# cat path_x, path_y to tensor of dim (B, 2*P)
|
||||
path = torch.cat([sample["path_x"], sample["path_y"]], dim=-1)
|
||||
path = self.path_net(path)
|
||||
path = self.path_net(sample["path"].reshape((sample["path"].shape[0], -1)))
|
||||
x = torch.cat([ego, relative, path], dim=-1)
|
||||
x = self.head(x)
|
||||
return x
|
||||
@@ -127,13 +124,13 @@ class IntersimPolicy():
|
||||
|
||||
def __call__(self, ob):
|
||||
|
||||
if 'action' in ob.keys():
|
||||
if 'ego_state' in ob.keys():
|
||||
# extract state from dataloader samples
|
||||
pass
|
||||
else:
|
||||
# extract state from observation (using simulator)
|
||||
ob['path_x'] = ob['paths'][0]
|
||||
ob['path_y'] = ob['paths'][1]
|
||||
ob['ego_state'] = ob['state']
|
||||
ob['path'] = torch.stack(ob['paths'],dim=-1)
|
||||
|
||||
ob = transform_observation(ob)
|
||||
|
||||
@@ -155,12 +152,14 @@ def generate_transforms(dataset):
|
||||
"""
|
||||
transforms = {
|
||||
'action': MinMaxScaler(),
|
||||
'state': MinMaxScaler(),
|
||||
'ego_state': MinMaxScaler(),
|
||||
'relative_state': MinMaxScaler(reduce_dim=2),
|
||||
'path_x': MinMaxScaler(reduce_dim=2),
|
||||
'path_y': MinMaxScaler(reduce_dim=2),
|
||||
'path': MinMaxScaler(reduce_dim=2),
|
||||
}
|
||||
for key in transforms.keys():
|
||||
if key == 'action':
|
||||
transforms[key].fit(dataset[:][key])
|
||||
else:
|
||||
transforms[key].fit(dataset[:]['state'][key])
|
||||
|
||||
return transforms
|
||||
|
||||
Reference in New Issue
Block a user