committing changes to start testing framework, removing shuffling of data

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

src/data/__init__.py

@@ -1 +1 @@
-from src.data.expert import single_agent_demonstrations, multi_agent_demonstrations, NoShuffleRNG, load_experts, process_experts 
+from src.data.expert import single_agent_expert, single_agent_demonstrations, multi_agent_demonstrations, NoShuffleRNG, load_experts, process_experts 

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', 
+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 <flags>
     Args: 
-        expert (class): class of expert
+        env (class): intersimple environment
-        env (class): class of env intersim.envs.intersimple
+        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:
@@ -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,