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

BUG FIXES: moving around when policy is loaded, adding BaseAlgorithm abstract classes, correcting metrics, normalizng actions if idm environment is a normalized action one, manually updating environment graph when using idm, implementing idm forward class

Browse Source
This commit is contained in:
Arec
2022-02-05 21:48:56 -08:00
parent 795e1c08b6
commit 3e6fce42ee
6 changed files with 195 additions and 197 deletions
Download Patch File Download Diff File Expand all files Collapse all files

16
evaluate_models.sh Executable file
View File

@@ -0,0 +1,16 @@
# eval_main inputs
# locations: List[Tuple[int,int]]= [(0,0)],
# method: str='expert',
# policy_file: str='',
# policy_kwargs: dict={},
# env: str='NRasterizedRouteIncrementingAgent',
# env_kwargs: dict={},
# seed: int=0
# expert
python -m src.eval_main
# idm
python -m src.eval_main --method=idm

1
src/__init__.py
View File

@@ -1,3 +1,2 @@
from src.data.expert_data import generate_expert_data, load_expert_data
from src.data.data_utils import InteractionDatasetSingleAgent
from src.evaluation.metrics import metrics

96
src/baselines/rule_policies.py
View File

@@ -1,6 +1,6 @@
from stable_baselines3.common.base_class import BaseAlgorithm
from intersim.envs.intersimple import Intersimple
from typing import Tuple, Optional
from intersim.envs.intersimple import Intersimple, NormalizedActionSpace
from typing import Tuple, Optional, List
import numpy as np
class PControllerPolicy(BaseAlgorithm):
@@ -10,11 +10,17 @@ class PControllerPolicy(BaseAlgorithm):
"""
Initialize policy with pointer to environment it will run on
"""
assert(isinstance(env, Intersimple), 'Environment is not an intersimple environment')
assert isinstance(env, Intersimple), 'Environment is not an intersimple environment'
self._env = env
self.target_v = 8.94 # m/s
self.attn_weight = 20
# BaseAlgorithm abstract methods
def _setup_model(self):
return None
def learn(self, *args, **kwargs):
return self
def predict(self, observation: np.ndarray, *args, **kwargs):
"""
Generate action, state from observation
@@ -59,37 +65,45 @@ class IDMRulePolicy(BaseAlgorithm):
"""
def __init__(self, env: Intersimple, target_speed: float= 8.94, t_future:float=0):
def __init__(self, env: Intersimple,
target_speed:float= 8.94,
t_future:List[float]=[0., 1., 2., 3.],
half_angle:float=60.):
"""
Initialize policy with pointer to environment it will run on and target speed
Args:
env (Intersimple): intersimple environment which IDM runs on
target_speed (float): target speed in roundabout (default: 8.94=20 mph)
t_future (float): future time at which to compare closest
t_future (List[float]): list of future time at which to compare closest vehicle
half_angle (float): half angle to look inside for closest vehicle
"""
# assert(isinstance(env, Intersimple), 'Environment is not an intersimple environment')
self._env = env
assert(t_future >=0, 'negative target speed')
self.t_future = t_future
self.half_angle = 45 # degrees for finding car to follow
self.half_angle = half_angle
# Default IDM parameters
assert(target_speed>0, 'negative target speed')
assert target_speed>0, 'negative target speed'
self.s_max = target_speed
self.a_max = np.array([3.]) # nominal acceleration
self.tau = 0.5 # desired time headway
self.b_pref = 2.5 # preferred deceleration
self.d_min = 1 #minimum spacing
super().__init__()
# for np.remainder nan warnings
np.seterr(invalid='ignore')
# BaseAlgorithm abstract methods
def _setup_model(self):
return None
def learn(self, *args, **kwargs):
return self
def predict(self, observation:np.ndarray,
*args, **kwargs) -> Tuple[np.ndarray,Optional[np.ndarray]]:
*args, **kwargs) -> Tuple[np.ndarray, None]:
"""
Generate action, state from observation
Predict action, state from observation
(But actually generate next action from underlying environment state)
@@ -98,41 +112,59 @@ class IDMRulePolicy(BaseAlgorithm):
Returns
action (np.ndarray): action for controlled agent to take
state (np.ndarray): the index of the chosen vehicle for IDM
state (None): None (hidden state for a recurrent policy)
"""
return self.forward(observation, *args, **kwargs), None
def forward(self, *args, **kwargs) -> np.ndarray:
"""
Generate action from underlying environment
Returns
action (np.ndarray): action for controlled agent to take
"""
import pdb
pdb.set_trace()
agent = self._env._agent
full_state = self._env._env.projected_state.numpy() #(nv, 5)
ego_state = full_state[agent] # (5,)
s = ego_state[2]
xy = full_state[:,0:2] # (nv, 2)
v = full_state[:,2:3] # (nv, 1)
psi = full_state[:,2:3] # (nv, 1)
psi = full_state[:,3:4] # (nv, 1)
d, r, i = self.get_ego_dr(agent, xy, v, psi)
# propagate environment forward at constant velocity
if self.t_future > 0:
xy2 = xy + self.t_future * v * np.vstack((np.cos(psi[:,0]), np.sin(psi[:,0])))
for t in self.t_future:
if t > 0:
xy2 = xy + t * v * np.vstack((np.cos(psi[:,0]), np.sin(psi[:,0]))).T
d2, r2, i2 = self.get_ego_dr(agent, xy2, v, psi)
# choose closer vehicle (now vs imagined)
if d2 < d:
d, r, i = d2, r2, i2
# Update environment interaction graph with i
if i:
self._env._env._graph._neighbor_dict={agent:[i]}
if d == np.inf:
d_des = 0
d_des = self.d_min
else:
d_des = self.d_min + self.tau * s + s * r / (2* (self.a_max*self.b_pref)**0.5 )
d_des = max(d_des, self.d_min)
assert (d_des>= self.d_min)
action = self.a_max*(1 - (s/self.s_max)**4 - (d_des/d)**2)
assert(action.shape==(1,))
return action, i
# normalize action to range if env is a NormalizedActionSpace
if isinstance(self._env, NormalizedActionSpace):
action = self._env._normalize(action)
assert action.shape==(1,)
return action
def get_ego_dr(self, agent:int, xy: np.ndarray,
v: np.ndarray, psi: np.ndarray) -> Tuple[Optional[np.ndarray], float, float]:
v: np.ndarray, psi: np.ndarray) -> Tuple[float, float, Optional[int]]:
"""
Return distance and relative speed of closest car within half angle from heading
@@ -145,14 +177,14 @@ class IDMRulePolicy(BaseAlgorithm):
Returns:
d (float): distance to closest vehicle in cone
r (float): relative speed between the two vehicles
i (Union[None,np.ndarray]): (1,) array of closest vehicle index, or None
i (Optional[int]): index of closest vehicle, or None
"""
nv, nxy = xy.shape
nv2, nvel = v.shape
nv3, npsi = psi.shape
assert(nv==nv2==nv3)
assert(nxy==2)
assert(nvel==npsi==1)
assert nv==nv2==nv3
assert nxy==2
assert nvel==npsi==1
dxys = xy - xy[agent] # (nv, 2)
ds = np.linalg.norm(dxys,axis=1) # (nv,)
@@ -160,7 +192,7 @@ class IDMRulePolicy(BaseAlgorithm):
dl = (dxys*np.hstack((-np.sin(psi), np.cos(psi)))).sum(-1) # (nv, )
alpha = to_circle(np.arctan2(dl, df))
val_idx = np.arange(nv)[(np.abs(alpha) < self.half_angle*np.pi/180 & np.arange(nv) != agent)]
val_idx = np.arange(nv)[(np.abs(alpha) < self.half_angle*np.pi/180) & (np.arange(nv) != agent)]
if len(val_idx)==0:
i = None
@@ -168,7 +200,7 @@ class IDMRulePolicy(BaseAlgorithm):
r = float('inf')
else:
idx = np.argmin(ds[val_idx]) # closest car which meets requirements
i = val_idx[idx]
i = int(val_idx[idx])
d = ds[i]
r = v[i,0]-v[agent,0]

160
src/eval_main.py
View File

@@ -2,18 +2,23 @@ from tqdm import tqdm
from copy import deepcopy
import stable_baselines3 as sb3
import intersim
from intersim.envs import Intersimple
from stable_baselines3.common.base_class import BaseAlgorithm
from src.baselines import IDMRulePolicy
from src.evaluation import IntersimpleEvaluation
import src.options.envs as options_envs
import src.gail.options as options_envs
from src.evaluation.metrics import divergence, visualize_distribution
from typing import Optional, List, Dict, Tuple
import torch
import numpy as np
def load_policy(policy_file:str, method:str,
policy_kwargs:dict, skip_load:bool=False) -> Optional[BaseAlgorithm]:
#(method, policy_file, policy_kwargs, eval_env)
def load_policy(method:str,
policy_file:str,
policy_kwargs:dict,
env: Intersimple) -> BaseAlgorithm:
"""
Load a model given a path and the method
@@ -21,13 +26,12 @@ def load_policy(policy_file:str, method:str,
load_policy (str): the path to the model
method (str): the method for the model
skip_load (bool): whether to skip loading
env (Intersimple): Intersimple environment for evaluation (necessary for IDM policy)
Returns:
policy (Optional[BaseAlgorithm]): the policy to evaluate
"""
if skip_load:
return None
elif method == 'idm':
policy = IDMRulePolicy(policy_kwargs)
if method == 'idm':
policy = IDMRulePolicy(env, **policy_kwargs)
elif method == 'bc':
raise NotImplementedError
elif method == 'gail':
@@ -67,9 +71,6 @@ def form_expert_metrics(states:torch.Tensor, actions:torch.Tensor) ->Dict[str, l
keys = ['col_all','v_all', 'a_all','j_all', 'v_avg', 'a_avg', 'col', 'brake', 't']
metrics = {key:[None]*nv for key in keys}
import pdb
pdb.set_trace()
for i in range(nv):
nni = ~torch.isnan(states[:,i,0])
@@ -115,7 +116,6 @@ def generate_expert_metrics(locations: List[Tuple[int,int]]) -> List[Dict[str, l
expert_metrics.append(form_expert_metrics(states, actions))
return expert_metrics
def load_expert_states(roundabout:int, track:int):
"""
Load expert states from roundabout/track info
@@ -127,16 +127,18 @@ def load_expert_states(roundabout:int, track:int):
actions (torch.tensor): (T, nv, 1) expert actions for track file
"""
rname = intersim.LOCATIONS[roundabout]
import pdb
pdb.set_trace()
state_path = 'expert_data/%s/track%04i/joint_expert_states.pt'%(rname, track)
action_path = 'expert_data/%s/track%04i/joint_expert_actions.pt'%(rname, track)
states = torch.load(state_path)
actions = torch.load(action_path)
return states, actions
def evaluate_policy(policy:BaseAlgorithm, locations:List[Tuple[int,int]],
env_class:str, env_kwargs:dict) -> List[Dict[str,list]]:
def evaluate_policy(locations:List[Tuple[int,int]],
env_class:str,
env_kwargs:dict,
method: str,
policy_file: str,
policy_kwargs:dict) -> List[Dict[str,list]]:
"""
Evaluate policy on an incrementing agent environment at all locations.
Return metrics for that policy
@@ -146,14 +148,16 @@ def evaluate_policy(policy:BaseAlgorithm, locations:List[Tuple[int,int]],
locations (list of tuples): list of locations to evaluate policy
env_class (str): name of environment to evaluate policy with
env_kwargs (dict): key word arguments to initialize environment with
method (str): policy method
policy_file (str): policy file path
policy_kwargs (dict): policy kwargs
Returns:
policy_metrics (list of dicts): policy_metrics[i][j][k] returns the value of metric 'j'
evaluated on the kth episode (car) of the ith roundabout trackfile under a policy
"""
envs_dict = deepcopy(intersim.envs.intersimple.__dict__)
envs_dict.update(deepcopy(options_envs.__dict__))
envs_dict = dict(intersim.envs.intersimple.__dict__)
envs_dict.update(dict(options_envs.__dict__))
policy_metrics = [None]* len(locations)
# iterate through vehicles
@@ -173,62 +177,71 @@ def evaluate_policy(policy:BaseAlgorithm, locations:List[Tuple[int,int]],
Env = envs_dict[env_class]
eval_env = Env(**env_kwargs)
evaluator = IntersimpleEvaluation(eval_env)
# load policy
policy = load_policy(method, policy_file, policy_kwargs, eval_env)
# run policy on environment
policy_metrics[i] = evaluator.evaluate(policy)
return policy_metrics
def summary_metrics(metrics:List[Dict[str,list]]):
def summary_metrics(metrics:List[Dict[str,list]]) -> Dict[str,float]:
"""
Summarize and print metrics averaged over vehicles and roundabouts
Args:
metrics (list of dicts): policy_metrics[i][j][k] returns the value of metric 'j'
evaluated on the kth episode (car) of the ith roundabout trackfile under a policy
Returns:
summary_metrics (Dict[str,float]): maps summary metric descriptions to values
"""
# keys = ['col_all','v_all', 'a_all','j_all', 'v_avg', 'a_avg', 'col', 'brake', 't']
import pdb
pdb.set_trace()
summary_metrics = {}
# average average-velocity
all_vavgs = sum([d['v_avg'] for d in metrics],[]) # aggregate to single list
mean_vavg = sum(all_vavgs)/len(all_vavgs)
print(f'Mean Average Velocity: {mean_vavg}')
summary_metrics['mean average velocity'] = sum(all_vavgs)/len(all_vavgs)
# average acceleration
all_aalls = np.concatenate([np.concatenate(d['a_all']) for d in metrics])
summary_metrics['mean acceleartion'] = np.mean(all_aalls)
# average +acceleration
pos_accels = all_aalls[all_aalls>0]
mean_pos_accels = np.mean(pos_accels)
print(f'Mean positive acceleration: {mean_pos_accels}')
summary_metrics['mean positive acceleration'] = np.mean(pos_accels)
# average deceleration
decels = all_aalls[all_aalls<0]
mean_decels = np.mean(decels)
print(f'Mean positive acceleration: {mean_decels}')
summary_metrics['mean deceleration'] = np.mean(decels)
# average jerk
all_jerks = np.concatenate([np.concatenate(d['j_all']) for d in metrics])
summary_metrics['mean jerk'] = np.mean(all_jerks)
# average |jerk|
all_jerks = np.concatenate([np.concatenate(d['j_all']) for d in metrics])
mean_abs_jerk = np.mean(np.abs(all_jerks))
print(f'Mean |Jerk|: {mean_abs_jerk}')
summary_metrics['mean |jerk|'] = np.mean(np.abs(all_jerks))
# collision rate
all_collisions = sum([d['col'] for d in metrics],[]) # aggregate to single list
collision_rate = sum(all_collisions)/len(all_collisions)
print(f'Collision Rate: {collision_rate}')
summary_metrics['collision rate'] = sum(all_collisions)/len(all_collisions)
# hard brake rate
all_hard_brakes = sum([d['brake'] for d in metrics],[]) # aggregate to single list
hard_brake_rate = sum(all_hard_brakes)/len(all_hard_brakes)
print(f'Hard Brake Rate: {hard_brake_rate}')
summary_metrics['hard brake rate'] = sum(all_hard_brakes)/len(all_hard_brakes)
# average number of timesteps
all_ts = sum([d['t'] for d in metrics],[]) # aggregate to single list
mean_t = sum(all_ts)/len(all_ts)
print(f'Mean episode length: {mean_t}')
summary_metrics['mean episode length'] = sum(all_ts)/len(all_ts)
def comparison_metrics(policy_metrics:List[Dict[str,list]], expert_metrics:List[Dict[str,list]]):
for key in summary_metrics.keys():
print(f'{key}: {summary_metrics[key]}')
return summary_metrics
def comparison_metrics(policy_metrics:List[Dict[str,list]],
expert_metrics:List[Dict[str,list]]) -> Dict[str,float]:
"""
Provide distributional comparison between different sets of metrics
@@ -238,39 +251,41 @@ def comparison_metrics(policy_metrics:List[Dict[str,list]], expert_metrics:List[
expert_metrics (list of dicts): expert_metrics[i][j][k] returns the value of metric 'j'
evaluated on the kth episode (car) of the ith expert roundabout trackfile
Returns:
comparison_metrics (Dict[str,float]): dict mapping comparison metric description to value
"""
import pdb
pdb.set_trace()
comparison_metrics = {}
# average velocity shortfall
expert_vavg = np.array(sum([d['v_avg'] for d in expert_metrics],[]))
policy_vavg = np.array(sum([d['v_avg'] for d in policy_metrics],[]))
assert(len(expert_vavg)==len(policy_vavg))
mean_shortfall = np.mean(expert_vavg - policy_vavg)
print(f'Mean shortfall velocity: {mean_shortfall}')
assert len(expert_vavg)==len(policy_vavg)
comparison_metrics['mean shortfall velocity'] = np.mean(expert_vavg - policy_vavg)
# velocity JSD
expert_vs = np.concatenate([np.concatenate(d['v_all']) for d in expert_metrics])
policy_vs = np.concatenate([np.concatenate(d['v_all']) for d in policy_metrics])
vel_div = divergence(expert_vs, policy_vs)
print(f'Velocity distribution divergence: {vel_div}')
visualize_distribution(expert_vs, policy_vs, 'velocity_jsd.png')
expert_vs = torch.tensor(np.concatenate([np.concatenate(d['v_all']) for d in expert_metrics]))
policy_vs = torch.tensor(np.concatenate([np.concatenate(d['v_all']) for d in policy_metrics]))
comparison_metrics['velocity distribution divergence'] = divergence(expert_vs, policy_vs)
visualize_distribution(expert_vs, policy_vs, 'velocity_jsd')
# acceleration JSD
expert_as = np.concatenate([np.concatenate(d['a_all']) for d in expert_metrics])
policy_as = np.concatenate([np.concatenate(d['a_all']) for d in policy_metrics])
accel_div = divergence(expert_as, policy_as)
print(f'Velocity distribution divergence: {accel_div}')
visualize_distribution(expert_as, policy_as, 'accel_jsd.png')
expert_as = torch.tensor(np.concatenate([np.concatenate(d['a_all']) for d in expert_metrics]))
policy_as = torch.tensor(np.concatenate([np.concatenate(d['a_all']) for d in policy_metrics]))
comparison_metrics['acceleration distribution divergence'] = divergence(expert_as, policy_as)
visualize_distribution(expert_as, policy_as, 'accel_jsd')
# jerk JSD
expert_jerks = np.concatenate([np.concatenate(d['j_all']) for d in expert_metrics])
policy_jerks = np.concatenate([np.concatenate(d['j_all']) for d in policy_metrics])
jerk_div = divergence(expert_jerks, policy_jerks)
print(f'Jerk distribution divergence: {jerk_div}')
visualize_distribution(expert_jerks, policy_jerks, 'jerk_jsd.png')
expert_jerks = torch.tensor(np.concatenate([np.concatenate(d['j_all']) for d in expert_metrics]))
policy_jerks = torch.tensor(np.concatenate([np.concatenate(d['j_all']) for d in policy_metrics]))
comparison_metrics['jerk distribution divergence'] = divergence(expert_jerks, policy_jerks)
visualize_distribution(expert_jerks, policy_jerks, 'jerk_jsd')
def test_model(
for key in comparison_metrics.keys():
print(f'{key}: {comparison_metrics[key]}')
return comparison_metrics
def eval_main(
locations: List[Tuple[int,int]]= [(0,0)],
method: str='expert',
policy_file: str='',
@@ -289,30 +304,27 @@ def test_model(
env (str): environment class
method (str): method (expert, bc, gail, rail, hgail, hrail)
"""
print(f'Evaluating {method} on {env}')
# set seed
np.random.seed(seed)
torch.manual_seed(seed)
# load policy
policy = load_policy(policy_file, method, policy_kwargs, skip_load=(method=='expert'))
# load expert metrics
expert_metrics = generate_expert_metrics(locations)
# if we have a policy
if policy:
# no comparison for expert
if method=='expert':
summary_metrics(expert_metrics)
# otherwise evaluate policy on roundabouts and generate metrics
else:
# evaluate it on the given roundabouts
policy_metrics = evaluate_policy(policy, locations, env, env_kwargs)
#
policy_metrics = evaluate_policy(locations, env, env_kwargs, method, policy_file, policy_kwargs)
summary_metrics(policy_metrics)
comparison_metrics(policy_metrics, expert_metrics)
else:
# if no policy, only generate summary metrics for the expert
summary_metrics(expert_metrics)
if __name__=='__main__':
import fire
fire.Fire()
fire.Fire(eval_main)

20
src/evaluation/evaluation.py
View File

@@ -1,7 +1,7 @@
import numpy as np
from stable_baselines3.common.vec_env import VecEnv
from stable_baselines3.common.evaluation import evaluate_policy
from intersim.envs.intersimple import Intersimple
from intersim.envs.intersimple import Intersimple, IncrementingAgent
from typing import Callable, Dict, Optional
import os
import pickle
@@ -19,21 +19,18 @@ class IntersimpleEvaluation:
- whether there was a collision
- whether there was a hard brake
"""
def __init__(self, eval_env, use_pbar:bool=True):
def __init__(self, eval_env:IncrementingAgent, use_pbar:bool=True):
"""
Initialize evaluation environment with an Intersimple IncrementingAgent environment
Args:
eval_env (Intersimple.IncrementingAgent): evaluation environment that increments agent upon reset
eval_env (IncrementingAgent): evaluation environment that increments agent upon reset
use_pbar (bool): whether to use a progress bar
"""
# if env is a VecEnv, the code needs to be adapted, since the callback will be called after each step,
# so transitions of different envs will be mixed and the total number of episodes could be larger than n_eval_episodes!
assert not isinstance(eval_env, VecEnv)
# make sure we have specified an environment which increments the agent number on reset
assert isinstance(eval_env, Intersimple.IncrementingAgent)
self.env = eval_env
self.n_episodes = eval_env.nv
self.use_pbar = use_pbar
@@ -54,7 +51,7 @@ class IntersimpleEvaluation:
"""
Reset metrics prior to evaluation
"""
self._metrics = {key: [[]]*self.n_episodes for key in self.metric_keys_all}
self._metrics = {key: [[] for _ in range(self.n_episodes)] for key in self.metric_keys_all}
self._metrics.update({key: [None]*self.n_episodes for key in self.metric_keys_single})
def save(self, filestr):
@@ -66,8 +63,8 @@ class IntersimpleEvaluation:
"""
# assert metrics all have correct length
for key in self.metric_keys:
assert(len(self._metrics[key])==self.n_episodes,
f'_metrics[{key}] does not have length {self.n_episodes}')
assert len(self._metrics[key])==self.n_episodes, \
f'_metrics[{key}] does not have length {self.n_episodes}'
# make filepath
os.makedirs(os.path.dirname(filestr))
@@ -91,7 +88,7 @@ class IntersimpleEvaluation:
evaluate_policy(
policy,
self.env,
n_eval_episodes=self.n_eval_episodes,
n_eval_episodes=self.n_episodes,
callback=self.evaluate_policy_callback,
return_episode_rewards=False
)
@@ -119,6 +116,7 @@ class IntersimpleEvaluation:
self._metrics['v_all'][_agent].append(info['prev_state'][_agent,2].item())
self._metrics['a_all'][_agent].append(info['action_taken'][_agent,0].item())
col = info['collision']
if col:
assert done
self._metrics['col_all'][_agent].append(col)
@@ -138,7 +136,7 @@ class IntersimpleEvaluation:
self._metrics['a_all'][i] = np.array(self._metrics['a_all'][i])
# jerk
self._metrics['j_all'][i] = np.diff(self._metrics['a_all'][i]) / self.eval_env._env._dt
self._metrics['j_all'][i] = np.diff(self._metrics['a_all'][i]) / self.env._env._dt
# average velocity and acceleration
self._metrics['v_avg'][i] = np.mean(self._metrics['v_all'][i])

79
src/evaluation/metrics.py
View File

@@ -4,83 +4,24 @@ import numpy as np
import matplotlib.pyplot as plt
from torch.utils.data import DataLoader
from intersim import collisions
# import tikzplotlib
def metrics(filestr: str, test_dataset, policy):
"""
Calculate metrics using a) base filestring to a simulation, and b) the test dataset and learned policy
Args:
filestr (str): base string to outputs of a simulation
test_dataset: a dataset held for testing
policy: policy
Returns:
info (dict): metrics in a dictionary
"""
info = {}
# compute metrics using either
# a) simulation files that were saved under the trained policy with prefix 'policy'
# b) applying the policy to observations in the test dataset
# load simulated trajectory
states = torch.load(filestr + '_sim_states.pt').detach()
lengths = torch.load(filestr + '_sim_lengths.pt').detach()
widths = torch.load(filestr + '_sim_widths.pt').detach()
xpoly = torch.load(filestr + '_sim_xpoly.pt').detach()
ypoly = torch.load(filestr + '_sim_ypoly.pt').detach()
# count collisions (from function in intersim.collisions)
n_collisions = collisions.count_collisions_trajectory(states, lengths, widths)
info['n_collisions'] = n_collisions
# calculate average velocity
avg_v = average_velocity(states)
info['average_velocity'] = avg_v
# convert policy dtype between float32 and float64
policy.policy = policy.policy.type(test_dataset[0]['state']['ego_state'].dtype)
# generate actions in test dataset
true_actions, pred_actions = [], []
true_velocities = []
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['state']))
true_actions.append(batch['action'])
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')
# calculate divergence between acceleration distributions
acceleration_divergence = divergence(pred_actions, true_actions, type='js')
info['acceleration_divergence'] = acceleration_divergence
# calculate divergence between velocity distributions
sim_velocities = states[:,:,2]
sim_velocities = sim_velocities[~torch.isnan(sim_velocities)].flatten()
true_velocities = torch.cat(true_velocities, dim=0)
velocity_divergence = divergence(sim_velocities, true_velocities, type='js')
info['velocity_divergence'] = velocity_divergence
return info
def visualize_distribution(true, pred, filestr):
def visualize_distribution(expert, policy, filestr):
"""
Visualize two distributions
Args:
true (torch.tensor): (n,)-sized true distribution
pred (torch.tensor): (m,)-sized pred distribution
expert (torch.tensor): (n,)-sized true distribution
generated (torch.tensor): (m,)-sized pred distribution
filestr (str): string to save figure to
"""
nni1 = ~torch.isnan(true)
nni2 = ~torch.isnan(pred)
nni1 = ~torch.isnan(expert)
nni2 = ~torch.isnan(policy)
plt.figure()
plt.hist(true[nni1].numpy(), density=True, bins=20)
plt.hist(pred[nni2].numpy(), density=True, bins=20)
plt.legend(['True', 'Predicted'])
plt.hist(expert[nni1].numpy(), density=True, bins=20)
plt.hist(policy[nni2].numpy(), density=True, bins=20)
plt.legend(['Expert', 'Predicted'])
plt.savefig(filestr+'.png')
# tikzplotlib.save(filestr+'.tex')
def average_velocity(states):
"""