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