diff --git a/scratch/arec/intersimple/test_model.py b/scratch/arec/intersimple/test_model.py index e4a5758..8ae5f2d 100644 --- a/scratch/arec/intersimple/test_model.py +++ b/scratch/arec/intersimple/test_model.py @@ -1,7 +1,9 @@ from tqdm import tqdm from copy import deepcopy +import stable_baselines3 as sb3 +import intersim -ALL_OPTIONS = +ALL_OPTIONS = [(v,t) for v in [0,2,4,6,8] for t in [5, 10]] # option 0 is safe fallback def load_model(model_path:str, method:str): """ @@ -17,7 +19,7 @@ def load_model(model_path:str, method:str): model = None is_heir = False if method == 'expert': - pass + raise NotImplementedError elif method == 'bc': raise NotImplementedError elif method == 'gail': @@ -26,7 +28,7 @@ def load_model(model_path:str, method:str): raise NotImplementedError elif method == 'hgail': is_heir = True - raise NotImplementedError + model = sb3.PPO.load(model_path) elif method == 'hrail': is_heir = True raise NotImplementedError @@ -41,12 +43,20 @@ def load_expert_states(roundabout, track): roundabout (str): roundabout name track (str): track id Returns: - expert_states (torch.tensor): (nv, T, 5) expert states for track file + states (torch.tensor): (T+1, nv, 5) expert states for track file + actions (torch.tensor): (T, nv, 1) expert actions for track file """ - pass + state_path = '../../../expert_data/%s/track%04i/joint_expert_states.pt'%(roundabout, track)] #FIXME when moving + action_path = '../../../expert_data/%s/track%04i/joint_expert_actions.pt'%(roundabout, track)] #FIXME when moving + states = torch.load(path) + actions = torch.load(path) + # nanify actions where vehicle's don't exist + import pdb + pdb.set_trace() + return states, actions def test_model( - locations=[], + locations=[(0,0)], model_name='gail_image_multiagent_nocollision', env='NRasterizedRouteIncrementingAgent', method='expert', @@ -56,7 +66,7 @@ def test_model( Test a particular model at different locations/tracks Args: - locations (list of tuples): list of (roundabout, track) pairs + locations (list of tuples): list of (roundabout, track) integer pairs model_name (str): name of model to test env (str): environment class method (str): method (expert, bc, gail, rail, hgail, hrail) @@ -70,23 +80,26 @@ def test_model( all_vehicle_infos = [] for i, location in tqdm(enumerate(locations)): - # load expert states - expert_states = load_expert_states(roundabout, track) - # add roundabout and track to environent roundabout, track = location + iround = intersim.LOCATIONS.index(roundabout) it_env_kwargs = deepcopy(env_kwargs) - it_env_kwargs.update({}) + loc_kwargs = { + 'loc':iround, + 'track':track + } + it_env_kwargs.update(loc_kwargs) # load expert states and get average velocities - expert_states = load_expert_states(roundabout, track) + expert_states, expert_actions = load_expert_states(roundabout, track) expert_vavg = torch.nanmean(expert_states[:,:,3], dim=-1) # initialize environment if not is_heir: - pass + Env = src.options.envs.__dict__[env] else: - pass + Env = intersim.envs.intersimple.__dict__[env] + env = Env(**env_kwargs) s = env.reset() # Iterate through every vehicle and time diff --git a/src/data/__init__.py b/src/data/__init__.py index e2f723c..f4dced8 100644 --- a/src/data/__init__.py +++ b/src/data/__init__.py @@ -1 +1 @@ -from src.data.expert import single_agent_demonstrations, multi_agent_demonstrations, NoShuffleRNG, load_experts, process_experts \ No newline at end of file +from src.data.expert import single_agent_expert, single_agent_demonstrations, multi_agent_demonstrations, NoShuffleRNG, load_experts, process_experts \ No newline at end of file diff --git a/src/data/expert.py b/src/data/expert.py index e06649b..0abd2c8 100644 --- a/src/data/expert.py +++ b/src/data/expert.py @@ -112,7 +112,7 @@ class NoShuffleRNG(np.random.RandomState): def shuffle(self, x): return x -def load_experts(expert_files, flatten = True): +def load_experts(expert_files, flatten=True): """ Load expert trajectories from files and combine their transitions into a single RB @@ -131,8 +131,27 @@ def load_experts(expert_files, flatten = True): transitions = rollout.flatten_trajectories(transitions) return transitions -def single_agent_demonstrations(expert='NormalizedIntersimpleExpert', - env='NRasterizedRouteIncrementingAgent', +def single_agent_expert(expert='NormalizedIntersimpleExpert', + env='NRasterizedRouteIncrementingAgent', + env_args={}, policy_args={}, **kwargs): + """ + Args: + expert (class): class of expert + env (class): class of env intersim.envs.intersimple + env_args (dict): dictionary of kwargs when instantiating environment class + policy_args (dict): dictionary of kwargs when instantiating Expert policy + path (str): path to store output + min_timesteps (int): min number of timesteps for call to rollout.rollout_and_save + min_episodes (int): min number of episodes for call to rollout.rollout_and_save + video (bool): whether to save a video of the expert until a single environment instantiation stops + """ + Env = intersim.envs.intersimple.__dict__[env] + Expert = globals()[expert] + env = Env(**env_args) + policy = Expert(env, **policy_args) + single_agent_demonstrations(env, policy, **kwargs) + +def single_agent_demonstrations(env, policy, path=None, min_timesteps=None, min_episodes=None, video=False, env_args={}, policy_args={}): @@ -141,8 +160,8 @@ def single_agent_demonstrations(expert='NormalizedIntersimpleExpert', Usage: python -m intersimple.expert Args: - expert (class): class of expert - env (class): class of env intersim.envs.intersimple + env (class): intersimple environment + policy (BasePolicy): intersimple policy path (str): path to store output min_timesteps (int): min number of timesteps for call to rollout.rollout_and_save min_episodes (int): min number of episodes for call to rollout.rollout_and_save @@ -151,14 +170,8 @@ def single_agent_demonstrations(expert='NormalizedIntersimpleExpert', policy_args (dict): dictionary of kwargs when instantiating Expert policy """ - Env = intersim.envs.intersimple.__dict__[env] - Expert = globals()[expert] - - env = Env(**env_args) info_env = RolloutInfoWrapper(env) # getting rollout info (dictionary) from environment venv = DummyVecEnv([lambda: info_env]) # making a DummyVecEnv with a list of a function that when called returns the rollout info - - policy = Expert(env, **policy_args) # instantiate an expert policy from specified class with instantiated environment and policy kwargs venv_policy = DummyVecEnvPolicy([lambda: policy]) # make a DummyVecEnvPolicy with a list of a function that when called returns the Expert policy if min_timesteps is None and min_episodes is None: @@ -166,7 +179,7 @@ def single_agent_demonstrations(expert='NormalizedIntersimpleExpert', if video: save_video(env, policy) - + path = path or (policy.__class__.__name__ + '_' + env.__class__.__name__ + '.pkl') suntil = rollout.make_sample_until( min_timesteps=min_timesteps, @@ -253,7 +266,7 @@ def process_experts(filename:str='expert.pkl', policy_args=expert_args ) # Single-Agent POV Demonstrations - single_agent_demonstrations( + single_agent_expert( expert=expert_class, env=env_class, path=it_path,