From 1a74fa5237f88479ca64faef8351a9542301e91a Mon Sep 17 00:00:00 2001 From: Arec Date: Tue, 20 Jul 2021 07:05:09 -0700 Subject: [PATCH 1/2] making general-purpose metric function --- src/main.py | 10 ++++------ src/metrics.py | 47 +++++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 51 insertions(+), 6 deletions(-) create mode 100644 src/metrics.py diff --git a/src/main.py b/src/main.py index c071dc4..b7068e5 100644 --- a/src/main.py +++ b/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): """ @@ -31,15 +31,13 @@ def main(method='bc', train=False, test=False, loc=0, **kwargs): 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 @@ -60,8 +58,8 @@ def main(method='bc', train=False, test=False, loc=0, **kwargs): 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=''): diff --git a/src/metrics.py b/src/metrics.py new file mode 100644 index 0000000..cfe0568 --- /dev/null +++ b/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 \ No newline at end of file From 6794b4cad8bc5851b2b527fd08ab4fd776aeacd7 Mon Sep 17 00:00:00 2001 From: Arec Date: Tue, 20 Jul 2021 08:29:00 -0700 Subject: [PATCH 2/2] making saving and loading functions class requirements, working on behavior cloning policy class and training function --- src/__init__.py | 3 +- src/bc/__init__.py | 2 +- src/bc/bc.py | 111 ++++++++++++++++++++++++++++++++++++++++++--- src/main.py | 6 +-- 4 files changed, 110 insertions(+), 12 deletions(-) diff --git a/src/__init__.py b/src/__init__.py index db75ef3..daabc89 100644 --- a/src/__init__.py +++ b/src/__init__.py @@ -1,2 +1,3 @@ from src.expert_data import generate_expert_data, load_expert_data -from src.data_utils import InteractionDatasetSingleAgent \ No newline at end of file +from src.data_utils import InteractionDatasetSingleAgent +from src.metrics import metrics \ No newline at end of file diff --git a/src/bc/__init__.py b/src/bc/__init__.py index 2b0a673..54f5c8d 100644 --- a/src/bc/__init__.py +++ b/src/bc/__init__.py @@ -1 +1 @@ -from src.bc.bc import BehaviorCloningPolicy, train, load_policy, metrics +from src.bc.bc import BehaviorCloningPolicy, train diff --git a/src/bc/bc.py b/src/bc/bc.py index 5b9a368..59d5ca0 100644 --- a/src/bc/bc.py +++ b/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) diff --git a/src/main.py b/src/main.py index b7068e5..63f240c 100644 --- a/src/main.py +++ b/src/main.py @@ -30,14 +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)) - - # method-based training if method=='bc': from src import bc policy_class = bc.BehaviorCloningPolicy - load_policy_fn = bc.load_policy train_fn = bc.train else: raise NotImplementedError @@ -51,7 +48,8 @@ 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