huangfu/InteractionImitation
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Merge branch 'main' of github.com:sisl/InteractionImitation

Browse Source
This commit is contained in:
Johannes Fischer
2021-08-04 19:44:04 +02:00
parent 5c40de66fa f9729b0a9d
commit 2224e2cd14
20 changed files with 3894 additions and 68 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
experiments/experiment.py
View File

@@ -48,6 +48,8 @@ def parse_args():
help='seed') help='seed')
parser.add_argument('--nframes', default=500, type=int, parser.add_argument('--nframes', default=500, type=int,
help='frames for test animation') 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', parser.add_argument('--graph', action='store_true',
help='whether to mask the relative states based on a ConeVisibilityGraph') help='whether to mask the relative states based on a ConeVisibilityGraph')
parser.add_argument('-d', default='./expert_data', type=str, parser.add_argument('-d', default='./expert_data', type=str,
@@ -64,6 +66,7 @@ def parse_args():
'seed':args.seed, 'seed':args.seed,
'ray':args.ray, 'ray':args.ray,
'nframes':args.nframes, 'nframes':args.nframes,
'nsamples':args.nsamples,
'datadir':os.path.abspath(args.d), 'datadir':os.path.abspath(args.d),
'graph':None, 'graph':None,
'outdir': opj('output',args.method,'loc%02i'%(args.loc)), 'outdir': opj('output',args.method,'loc%02i'%(args.loc)),
@@ -156,7 +159,7 @@ if __name__ == '__main__':
local_dir=kwargs['outdir'], local_dir=kwargs['outdir'],
#resources_per_trial={"cpu": 2}, #resources_per_trial={"cpu": 2},
time_budget_s=120*60, time_budget_s=120*60,
num_samples=200, num_samples=kwargs['nsamples'],
) )
elif kwargs['ray'] and kwargs['test']: elif kwargs['ray'] and kwargs['test']:
analysis = Analysis(kwargs['outdir'], default_metric="cv_loss", default_mode="min") analysis = Analysis(kwargs['outdir'], default_metric="cv_loss", default_mode="min")

274
scratch/etienne/pillbox/intersim_advil.ipynb Normal file
View File

File diff suppressed because one or more lines are too long

249
scratch/etienne/pillbox/intersim_demos.ipynb Normal file
View 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
View 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
View File

File diff suppressed because one or more lines are too long

140
scratch/etienne/pillbox/learners/adril.py Normal file
View 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
View 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
View 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
View 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
View 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
View 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
View 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
View File

@@ -0,0 +1,9 @@
gym
numpy
psutil
scikit_learn
scipy
stable_baselines3
torch
tqdm
imitation

0
scratch/etienne/test.py
View File

6
src/bc/bc.py
View File

@@ -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) cv_loader = DataLoader(cv_dataset, batch_size=cv_batch_size, shuffle=True)
# change policy dtype # 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 # generate loss function, optimizer
cv_loss_fn = nn.MSELoss(reduction='sum') 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): for (batch_idx, batch) in enumerate(training_loader):
# sample mini-batch and run through policy # sample mini-batch and run through policy
pred_action = policy(batch) pred_action = policy(batch['state'])
loss = loss_fn(pred_action, batch['action']) loss = loss_fn(pred_action, batch['action'])
# compute loss and step optimizer # compute loss and step optimizer
@@ -169,7 +169,7 @@ def train(config, policy, train_dataset, cv_dataset, filestr, **kwargs):
with torch.no_grad(): with torch.no_grad():
cv_loss = 0. cv_loss = 0.
for (batch_idx, batch) in enumerate(cv_loader): 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']) loss = cv_loss_fn(pred_action, batch['action'])
cv_loss += loss.item() / len(cv_dataset) cv_loss += loss.item() / len(cv_dataset)

57
src/data_utils.py
View File

@@ -25,13 +25,14 @@ class InteractionDatasetSingleAgent(Dataset):
self.loc = loc self.loc = loc
self.tracks = tracks self.tracks = tracks
self.dtype = dtype self.dtype = dtype
self.keys = ['ego_state', 'relative_state', 'path', 'action', 'next_ego_state', 'next_relative_state', 'next_path']
self._load_dataset() self._load_dataset()
def _load_dataset(self): def _load_dataset(self):
""" """
Load the full datasets ahead of time 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 max_nv = 0
for track in self.tracks: for track in self.tracks:
try: try:
@@ -41,33 +42,26 @@ class InteractionDatasetSingleAgent(Dataset):
print('Failed to load location {} track {}'.format(self.loc,track)) print('Failed to load location {} track {}'.format(self.loc,track))
continue continue
max_nv = max(max_nv, data['relative_state'].shape[1]) max_nv = max(max_nv, data['relative_state'].shape[1])
self.raw_data['state'].append(data['state']) for key in self.keys:
self.raw_data['relative_state'].append(data['relative_state']) self.raw_data[key].append(data[key])
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)
# pad second dimension of relative state # pad second dimension of relative state
for i in range(len(self.raw_data['relative_state'])): for i in range(len(self.raw_data['relative_state'])):
nv1, nv2, d = self.raw_data['relative_state'][i].shape nv1, nv2, d = self.raw_data['relative_state'][i].shape
pad = torch.zeros(nv1, max_nv-nv2, d, dtype=self.dtype) * np.nan 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'][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 # mandate equal length
assert len(self.raw_data['state']) == len(self.raw_data['relative_state']) \ lengths = [len(self.raw_data[key]) for key in self.keys]
== len(self.raw_data['action']) \ assert min(lengths) == max(lengths), 'dataset lengths unequal'
== len(self.raw_data['path_x']) \
== len(self.raw_data['path_y']), 'dataset lengths unequal'
def __len__(self): def __len__(self):
return len(self.raw_data['state']) return len(self.raw_data['ego_state'])
def __getitem__(self, idx): def __getitem__(self, idx):
""" """
@@ -76,12 +70,27 @@ class InteractionDatasetSingleAgent(Dataset):
idx: index or indices of B samples idx: index or indices of B samples
Returns: Returns:
sample (dict): sample dictionary with the following entries: sample (dict): sample dictionary with the following entries:
state (torch.tensor): (B, 5) raw state state (dict): state dictionary with the following entries:
relative_state (torch.tensor): (B, max_nv, d) relative state (padded with nans) ego_state (torch.tensor): (B, 5) raw state
path_x (torch.tensor): (B, P) tensor of P future path x positions relative_state (torch.tensor): (B, max_nv, d) relative state (padded with nans)
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 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 self.keys}
sample = {key:self.raw_data[key][idx] for key in 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 return sample

49
src/expert_data.py
View File

@@ -31,8 +31,8 @@ def generate_expert_data(path: str='expert_data', loc: int = 0, track:int = 0,
os.makedirs(path) os.makedirs(path)
filestr = opj(path,intersim.LOCATIONS[loc]+'_track%03i'%(track)) filestr = opj(path,intersim.LOCATIONS[loc]+'_track%03i'%(track))
svt, svt_path = get_svt(base='InteractionSimulator', loc=loc, track=track) svt, svt_path = get_svt(loc=loc, track=track) #base='InteractionSimulator'
osm = get_map_path(base='InteractionSimulator', loc=loc) osm = get_map_path(loc=loc)
print('SVT path: {}'.format(svt_path)) print('SVT path: {}'.format(svt_path))
print('Map path: {}'.format(osm)) print('Map path: {}'.format(osm))
states, actions = SVT_to_stateactions(svt) 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 remove_outliers (bool): whether to remove datapoints with acceleration above or below 5 m/s/s
dtype (torch.Type): type to convert data to 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} data = {key:[] for key in keys}
assert len(obs) == len(actions), 'non-matching action and observation lengths' assert len(obs) == len(actions), 'non-matching action and observation lengths'
T = len(obs) T = len(obs)
max_nv = 0 max_nv = 0
for t in range(T): for t in range(T-1):
nni = ~torch.isnan(obs[t]['state'][:,0]) nni = ~torch.isnan(obs[t]['state'][:,0]) & ~torch.isnan(obs[t+1]['state'][:,0])
max_nv = max(max_nv,nni.count_nonzero()) 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, data['relative_state'].append(obs[t]['relative_state'].index_select(0,
nni.nonzero()[:,0]).index_select(1, nni.nonzero()[:,0])) nni.nonzero()[:,0]).index_select(1, nni.nonzero()[:,0]))
data['action'].append(actions[t][nni]) data['path'].append(torch.stack((obs[t]['paths'][0][nni], obs[t]['paths'][1][nni]), dim=-1))
data['path_x'].append(obs[t]['paths'][0][nni])
data['path_y'].append(obs[t]['paths'][1][nni]) # 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 # pad second dimension of relative state
for i in range(len(data['relative_state'])): for i in range(len(data['relative_state'])):
nv1, nv2, d = data['relative_state'][i].shape nv1, nv2, d = data['relative_state'][i].shape
pad = torch.zeros(nv1, max_nv-nv2, d, dtype=dtype) * np.nan 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'][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: if remove_outliers:
non_outlier_indices = torch.nonzero(torch.abs(data['action'][:,0]) < 5) 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]] data[key] = data[key][non_outlier_indices[:,0]]
# mandate equal length # mandate equal length
assert len(data['state']) == len(data['relative_state']) \ lengths = [len(data[key]) for key in keys]
== len(data['action']) == len(data['path_x']) \ assert min(lengths) == max(lengths), 'dataset lengths unequal'
== len(data['path_y']), 'dataset lengths unequal'
# save out data # save out data
for key in data.keys(): for key in keys:
torch.save(data[key], filestr+'_'+key+'.pt') torch.save(data[key], filestr+'_'+key+'.pt')
def load_expert_data(path='expert_data', loc: int = 0, track:int = 0): 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 # load observations and actions
filestr = opj(path, intersim.LOCATIONS[loc]+'_track%03i'%(track)) filestr = opj(path, intersim.LOCATIONS[loc]+'_track%03i'%(track))
data = {} 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') data[key] = torch.load(filestr+'_'+key+'.pt')
return data return data

9
src/main.py
View File

@@ -47,11 +47,11 @@ def main(config, method='bc', train=False, test=False, loc=0, datadir='./expert_
if train: if train:
# make policy, train and test datasets, and send to # 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 # generate transform from train_dataset
transforms = generate_transforms(train_dataset) transforms = generate_transforms(train_dataset)
policy = policy_class(config, transforms) 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) train_fn(config, policy, train_dataset, cv_dataset, filestr, **kwargs)
if test: if test:
@@ -61,11 +61,10 @@ def main(config, method='bc', train=False, test=False, loc=0, datadir='./expert_
policy.eval() policy.eval()
# simulate policy # simulate policy
track = 4 simulate_policy(policy, loc=loc, track=kwargs['test_tracks'][0], filestr=filestr, nframes=kwargs['nframes'], graph=kwargs['graph'])
simulate_policy(policy, loc=loc, track=track, filestr=filestr, nframes=kwargs['nframes'], graph=kwargs['graph'])
# run test metrics # 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) writer = SummaryWriter(filestr)
info = metrics(filestr, test_dataset, policy) info = metrics(filestr, test_dataset, policy)
for k, m in info.items(): for k, m in info.items():

6
src/metrics.py
View File

@@ -37,7 +37,7 @@ def metrics(filestr: str, test_dataset, policy):
info['average_velocity'] = avg_v info['average_velocity'] = avg_v
# convert policy dtype between float32 and float64 # 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 # generate actions in test dataset
true_actions, pred_actions = [], [] true_actions, pred_actions = [], []
@@ -45,9 +45,9 @@ def metrics(filestr: str, test_dataset, policy):
test_loader = DataLoader(test_dataset, batch_size=1024) test_loader = DataLoader(test_dataset, batch_size=1024)
with torch.no_grad(): with torch.no_grad():
for (batch_idx, batch) in enumerate(test_loader): 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_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) 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') visualize_distribution(true_actions[:,0], pred_actions[:,0], filestr+'_action_viz')

23
src/policies/policy.py
View File

@@ -31,17 +31,14 @@ class IntersimStateNet(nn.Module):
sample (dict): sample dictionary with the following entries: sample (dict): sample dictionary with the following entries:
state (torch.tensor): (B, 5) raw state state (torch.tensor): (B, 5) raw state
relative_state (torch.tensor): (B, max_nv, d) relative state (padded with nans) 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 (torch.tensor): (B, P, 2) tensor of P future path x and y positions
path_y (torch.tensor): (B, P) tensor of P future path y positions
action (torch.tensor): (B, 1) actions taken from each state action (torch.tensor): (B, 1) actions taken from each state
Returns: Returns:
x (torch.tensor): (head_output_dim,) output of common head network x (torch.tensor): (head_output_dim,) output of common head network
""" """
ego = self.ego_net(sample["ego_state"]) ego = self.ego_net(sample["ego_state"])
relative = self.deepsets_net(sample["relative_state"]) relative = self.deepsets_net(sample["relative_state"])
# cat path_x, path_y to tensor of dim (B, 2*P) path = self.path_net(sample["path"].reshape((sample["path"].shape[0], -1)))
path = torch.cat([sample["path_x"], sample["path_y"]], dim=-1)
path = self.path_net(path)
x = torch.cat([ego, relative, path], dim=-1) x = torch.cat([ego, relative, path], dim=-1)
x = self.head(x) x = self.head(x)
return x return x
@@ -127,13 +124,13 @@ class IntersimPolicy():
def __call__(self, ob): def __call__(self, ob):
if 'action' in ob.keys(): if 'ego_state' in ob.keys():
# extract state from dataloader samples # extract state from dataloader samples
pass pass
else: else:
# extract state from observation (using simulator) # extract state from observation (using simulator)
ob['path_x'] = ob['paths'][0] ob['ego_state'] = ob['state']
ob['path_y'] = ob['paths'][1] ob['path'] = torch.stack(ob['paths'],dim=-1)
ob = transform_observation(ob) ob = transform_observation(ob)
@@ -155,12 +152,14 @@ def generate_transforms(dataset):
""" """
transforms = { transforms = {
'action': MinMaxScaler(), 'action': MinMaxScaler(),
'state': MinMaxScaler(), 'ego_state': MinMaxScaler(),
'relative_state': MinMaxScaler(reduce_dim=2), 'relative_state': MinMaxScaler(reduce_dim=2),
'path_x': MinMaxScaler(reduce_dim=2), 'path': MinMaxScaler(reduce_dim=2),
'path_y': MinMaxScaler(reduce_dim=2),
} }
for key in transforms.keys(): for key in transforms.keys():
transforms[key].fit(dataset[:][key]) if key == 'action':
transforms[key].fit(dataset[:][key])
else:
transforms[key].fit(dataset[:]['state'][key])
return transforms return transforms