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

src/__init__.py

@@ -1 +1,2 @@
-from src.expert_data import generate_expert_data, load_expert_data
+from src.expert_data import generate_expert_data, load_expert_data
+from src.data_utils import InteractionDatasetSingleAgent

src/bc/__init__.py

@@ -0,0 +1 @@
+from src.bc.bc import BehaviorCloningPolicy, train, load_policy, metrics

src/bc/bc.py

@@ -0,0 +1,13 @@
+class BehaviorCloningPolicy():
+    pass
+
+def load_policy():
+    pass
+
+def metrics():
+    pass
+
+def train():
+    pass
+
+

src/data_utils.py

@@ -1,7 +1,6 @@
 import torch
-from torch.utils.data import Dataset, DataLoader
+from torch.utils.data import Dataset
 import numpy as np
-#from torchvision import transforms, utils
 from src.expert_data import load_expert_data
 import os
 opj = os.path.join
@@ -15,24 +14,16 @@ class InteractionDatasetMultiAgent(Dataset):
 class InteractionDatasetSingleAgent(Dataset):
     """Class to load states and actions for individual agents."""
 
-    def __init__(self, output_dir='expert_data', loc:int = 0, tracks:list = [0], transforms={}):
+    def __init__(self, output_dir='expert_data', loc:int = 0, tracks:list = [0]):
         """
         Args:
             output_dir (string): Directory with all the images.
             loc (int): location index
             tracks (list[int]): track indices
-            transforms (dict): dictionary of transforms to apply to different variables
         """
         self.output_dir = output_dir
         self.loc = loc
         self.tracks = tracks
-        self.transforms = transforms
-        #self.action_transform = transforms.get('action', None)
-        #self.state_transform = transforms.get('state', None)
-        #self.relative_state_transform = transforms.get('relative_state', None)
-        #self.paths_x_transform = transforms.get('paths_x', None)
-        #self.paths_y_transform = transform.get('paths_y',None)
-
         self._load_dataset()
 
     def _load_dataset(self):
@@ -95,9 +86,4 @@ class InteractionDatasetSingleAgent(Dataset):
         """
         keys = ['state', 'relative_state', 'path_x', 'path_y', 'action']
         sample = {key:self.raw_data[key][idx] for key in keys}
-
-        for key in keys: 
-            if key in self.transforms.keys():
-                sample[key] = self.transforms[key](sample[key])
-
         return sample

src/expert_data.py

@@ -91,5 +91,11 @@ if __name__ == '__main__':
                        help='location (default 0)')
     parser.add_argument('--track', default=0, type=int,
                        help='track number (default 0)')
+    parser.add_argument('--all-tracks', action='store_true',
+                       help='whether to process all tracks at location')
     args = parser.parse_args()
-    generate_expert_data(loc=args.loc,track=args.track)
+    if args.all_tracks:
+        for i in range(intersim.MAX_TRACKS):
+            generate_expert_data(loc=args.loc, track=i)
+    else:
+        generate_expert_data(loc=args.loc,track=args.track)

src/main.py

@@ -0,0 +1,127 @@
+import torch
+import gym
+import intersim
+import numpy as np
+import os
+opj = os.path.join
+from src import InteractionDatasetSingleAgent
+
+def basestr(**kwargs):
+    """
+    Return base prefix for all files relating to a certain experiment
+    Args:
+        kwargs (dict): keyword arguments sent to main training loop
+    Returns:
+        basestr (str): prefix
+    """
+    return 'base_'
+
+def main(method='bc', train=False, test=False, loc=0, **kwargs):
+    """
+    Main loop for training and testing different imitation models
+    Args:
+        train (bool): whether to run train loop
+        test (bool): whether to run test loop
+        method (str): the method to try for imitation
+        loc (int): the location index of the roundabout
+        kwargs (dict): remaining kwargs for policy and training loop
+    """
+    outdir = opj('output',method,'loc%02i'%(loc))
+    if not os.path.isdir(outdir):
+        os.mkdir(outdir)
+    filestr = opj(outdir, basestr(**kwargs))
+   
+    # define transforms
+    transforms={}
+    
+    # method-based training
+    if method=='bc':
+        from src import bc
+        policy_class = bc.BehaviorCloningPolicy
+        load_policy_fn = bc.load_policy
+        metrics_fn = bc.metrics
+        train_fn = bc.train
+    else:
+        raise NotImplementedError
+    
+    # default train / cv / test split datasets 
+    if train:
+        train_dataset = InteractionDatasetSingleAgent(loc=loc, tracks=[0,1,2])
+        cv_dataset = InteractionDatasetSingleAgent(loc=loc, tracks=[3])
+        train_fn(train_dataset, cv_dataset, policy_class, filestr=filestr, **kwargs)
+    
+    if test:
+
+        # load policy 
+        policy = load_policy_fn(filestr=filestr)
+
+        # simulate policy
+        test_track = 4
+        simulate_policy(policy, loc=loc, track=track, filestr=filestr)
+
+        # run test metrics
+        # test_dataset = InteractionDatasetSingleAgent(loc=loc, tracks=[4])
+        # metrics_fn(test_dataset, policy)
+
+
+def simulate_policy(policy, loc=0, track=0, filestr=''):
+    """
+    Simulate a trained policy
+    Args:
+        policy: the policy to simulate, which should return action directly
+        loc (int): location index to test policy
+        track (int): track to test policy
+        filestr (str): path prefix to save simulation to
+    """
+    # animate from environment
+    env = gym.make('intersim:intersim-v0', loc=loc, track=track, 
+        min_acc=-np.inf, max_acc=np.inf)
+    
+    ob, _ = env.reset()
+    env.render()
+    done = False
+    while not done: 
+        # get action
+        action = policy(ob)
+
+        # propagate environment
+        ob, r, done, info = env.step(action)
+        env.render()
+
+    env.close(filestr=filestr)
+
+def parse_args():
+    """
+    Parse arguments to main
+    Returns:
+        kwargs: dictionary of arguments:
+            train (bool): whether to run train loop
+            test (bool): whether to run test loop
+            method (str): the method to try for imitation
+            loc (int): the location index of the roundabout
+    """
+    import argparse
+    parser = argparse.ArgumentParser(description='Save Expert Trajectories')
+    parser.add_argument('--loc', default=0, type=int,
+        help='location (default 0)')
+    parser.add_argument("--train", help="train model",
+        action="store_true")
+    parser.add_argument("--test", help="test model",
+        action="store_true") 
+    parser.add_argument("--method", help="modeling method",
+        choices=['bc', 'gail', 'advil'], default='bc')        
+    parser.add_argument()              
+    parser.add_argument()
+    args = parser.parse_args()
+    kwargs = {
+        'train'=args.train, 
+        'test'=args.test, 
+        'method'=args.method,
+        'loc'=args.loc
+    }
+    return kwargs
+
+
+if __name__ == '__main__':
+    kwargs = parse_args()
+    main(**kwargs)