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

src/__init__.py

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

src/bc/__init__.py

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

src/bc/bc.py

@@ -1,13 +1,112 @@
+import torch
+import torch.nn as nn
+from torch.utils.data import DataLoader
+import pickle
+
+from src.policies import 
+
 class BehaviorCloningPolicy():
+    """
+    Class for (continuous) behavior cloning policy
+    """
+    
+    def __init__(self, transforms={}, **kwargs):
+        self._transforms = transforms
+        self._policy_model = PolicyModel(**kwargs)
+
+    @property
+    def transforms(self):
+        return self._transforms
+    
+    @transforms.setter
+    def transforms(self, transforms)
+        self._transforms=transforms
+
+    def __call__(self, ob):
+
+        if 'actions' in ob.keys():
+            # extract state from dataloader samples  
+            pass
+        else:
+            # extract state from observation (using simulator)
+            ob['path_x'] = ob['paths'][0]
+            ob['path_y'] = ob['paths'][1]
+
+        # run observation through transforms
+        for key in ['state', 'relative_state', 'path_x', 'path_y']:             
+            if key in self._transforms.keys():
+                ob[key] = self._transforms[key](ob[key])
+
+        # run transformed state through model
+
+        # untransform action
+
+        pass
+
+    @classmethod
+    def load_model(cls, filestr, **kwargs):
+        """
+        Load a model from a file prefix
+        Args:
+            filestr (str): string prefix to load model from
+        Returns
+            model (BehaviorCloningPolicy): loaded model
+        """
+        transforms = pickle.load(open(filestr+'_transforms.pkl', 'rb'))
+        model = cls(transforms=transforms, **kwargs)
+        model._policy_model.load_state_dict(torch.load(filestr+'_model.pt'))
+        return model
+    
+    def eval(self):
+        self._policy_model.eval()
+
+    def save_model(self, filestr):
+        """
+        Save transforms and state_dict to a location specificed by filestr
+        Args:
+            filestr (str): string prefix to save model to
+        """
+        pickle.dump(self._transforms, open(filestr+'_transforms.pkl', 'wb'))
+        torch.save(self._policy_model.state_dict(), filestr+'_model.pt')
+
+def generate_transforms(dataset):
+    """
+    Generate transform dictionary from dataset
+    """
     pass
 
-def load_policy():
-    pass
+def train(train_dataset, cv_dataset, policy_class, filestr=filestr, **kwargs):
+    
+    # hyperparams
+    train_epochs = 10000
+    cv_every = 100
+    train_batch_size = 64
+    cv_batch_size = 256
 
-def metrics():
-    pass
+    # generate transform from train_dataset
+    transforms = generate_transforms(train_dataset)
 
-def train():
-    pass
+    # initialize policy
+    policy = BehaviorCloningPolicy(transforms=transforms, **kwargs)
+
+    # training and testing dataloaders
+    training_loader = DataLoader(train_dataset, batch_size=train_batch_size, shuffle=True)
+    cv_loader = DataLoader(cv_dataset, batch_size=cv_batch_size, shuffle=True)
+
+    # generate loss function, optimizer
+
+    for i in train_epochs:
+
+        # sample mini-batch and run through policy
+        pred_action = policy(batch) 
+
+        # compute loss and step optimizer
+        loss.backwards()
+        optimizer.step()
 
 
+        # measure L2 on every cv epoch
+        if i % cv_every == 0:
+            pass
+
+    policy.save_model(filestr)

src/main.py

@@ -4,7 +4,7 @@ import intersim
 import numpy as np
 import os
 opj = os.path.join
-from src import InteractionDatasetSingleAgent
+from src import InteractionDatasetSingleAgent, metrics
 
 def basestr(**kwargs):
     """
@@ -30,16 +30,11 @@ def main(method='bc', train=False, test=False, loc=0, **kwargs):
     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
@@ -53,15 +48,16 @@ def main(method='bc', train=False, test=False, loc=0, **kwargs):
     if test:
 
         # load policy 
-        policy = load_policy_fn(filestr=filestr)
+        policy = policy_class.load_model(filestr=filestr, **kwargs)
+        policy.eval()
 
         # 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)
+        test_dataset = InteractionDatasetSingleAgent(loc=loc, tracks=[4])
+        metrics(filestr, test_dataset, policy)
 
 
 def simulate_policy(policy, loc=0, track=0, filestr=''):

src/metrics.py

@@ -0,0 +1,47 @@
+import torch
+import pickle
+import numpy as np
+
+import intersim.collisions
+
+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
+    """
+
+    # 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 trajectory
+
+    # count collisions (from function in intersim.collisions)
+
+    # calculate average velocity
+
+    # calculate divergence between velocity distributions
+    
+    # calcuate divergence between acceleration distributions
+     
+    pass
+
+def average_velocity(x):
+    """
+
+    """
+    pass
+
+def divergence(p, q, type='kl')
+    """
+    Calculate a divergence between p and q
+    Args:
+        p (torch.tensor): (n) samples from p
+        q (torch.tensor): (n) samples from q
+    Returns:
+        d (float): approximate divergence
+    """
+    pass