diff --git a/src/bc/bc.py b/src/bc/bc.py index 59d5ca0..cd32297 100644 --- a/src/bc/bc.py +++ b/src/bc/bc.py @@ -1,19 +1,29 @@ import torch import torch.nn as nn -from torch.utils.data import DataLoader +from torch.utils.data import DataLoader, RandomSampler import pickle +from torch.utils.tensorboard import SummaryWriter -from src.policies import +from src.policies import DeepSetsPolicy +from src.util.transform import SciKitMinMaxScaler +import json5 class BehaviorCloningPolicy(): """ Class for (continuous) behavior cloning policy """ - def __init__(self, transforms={}, **kwargs): + def __init__(self, config: dict transforms: dict={}): + """ + Initialize BehaviorCloningPolicy + Args: + config (dict): configuration file to initialize DeepSetsPolicy with + transforms (dict): dictionary of transforms to apply to different fields + """ + self._config = config self._transforms = transforms - self._policy_model = PolicyModel(**kwargs) - + self._policy_model = DeepSetsPolicy(config["ego_state"], config["deepsets"], config["path_encoder"], config["head"]) + @property def transforms(self): return self._transforms @@ -35,31 +45,38 @@ class BehaviorCloningPolicy(): # 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]) + ob[key] = self._transforms[key].transform(ob[key]) # run transformed state through model + action = self._policy_model(ob) + assert action.ndim == 2, 'action has incorrect shape' # untransform action - - pass + if 'action' in self._transforms.keys(): + action = self._transforms['action'].inverse_transform(action) + return action @classmethod - def load_model(cls, filestr, **kwargs): + def load_model(cls, config: dict, filestr: str): """ Load a model from a file prefix Args: + config (dict): configuration dict to set up model 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 = cls(config, transforms=transforms) model._policy_model.load_state_dict(torch.load(filestr+'_model.pt')) return model def eval(self): self._policy_model.eval() + def parameters(self): + return self._policy_model.parameters() + def save_model(self, filestr): """ Save transforms and state_dict to a location specificed by filestr @@ -72,41 +89,72 @@ class BehaviorCloningPolicy(): def generate_transforms(dataset): """ Generate transform dictionary from dataset + Args: + dataset (Dataset): dataset of demo observations and actions """ - pass + transforms = { + 'action': SciKitMinMaxScaler() + 'state': SciKitMinMaxScaler() + 'relative_state': SciKitMinMaxScaler(reduce_dim=2) + 'path_x': SciKitMinMaxScaler(reduce_dim=2) + 'path_y': SciKitMinMaxScaler(reduce_dim=2) + } + for key in transforms.keys(): + transforms[key].fit(dataset[:][key]) -def train(train_dataset, cv_dataset, policy_class, filestr=filestr, **kwargs): + return transforms + +def train(train_dataset, cv_dataset, policy, filestr, **kwargs): # hyperparams train_epochs = 10000 cv_every = 100 train_batch_size = 64 - cv_batch_size = 256 + cv_batch_size = 256 # doesn't matter + learning_rate = 1e-3 + weight_decay=0.1 # generate transform from train_dataset transforms = generate_transforms(train_dataset) # initialize policy - policy = BehaviorCloningPolicy(transforms=transforms, **kwargs) + policy.transforms = transforms # 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 + loss_fn = nn.HuberLoss(reduction='sum') + optimizer = torch.optim.Adam(policy.parameters(), lr=learning_rate, weight_decay=weight_decay) for i in train_epochs: - # sample mini-batch and run through policy - pred_action = policy(batch) + epoch_loss = 0 + for (batch_idx, batch) in enumerate(training_loader): + # sample mini-batch and run through policy + pred_action = policy(batch) + loss = loss_fn(pred_action, batch['action']) - # compute loss and step optimizer - loss.backwards() - optimizer.step() + # compute loss and step optimizer + optimizer.zero_grad() + loss.backwards() + optimizer.step() + epoch_loss += loss.item() / len(train_dataset) + + # Write epoch loss + if i % 10 == 0: + print('Epoch: {}, Training Loss: {}'.format(i, epoch_loss)) - # measure L2 on every cv epoch + # measure cv loss if i % cv_every == 0: - pass - + with torch.no_grad(): + cv_loss = 0. + for (batch_idx, batch) in enumerate(cv_loader): + pred_action = policy(batch) + loss = loss_fn(pred_action, batch['action']) + cv_loss += loss.item() / len(cv_dataset) + print('Epoch: {}, CV Loss: {}'.format(i, cv_loss)) + policy.save_model(filestr) diff --git a/src/main.py b/src/main.py index 63f240c..cd7f8f5 100644 --- a/src/main.py +++ b/src/main.py @@ -2,8 +2,11 @@ import torch import gym import intersim import numpy as np +import json5 import os opj = os.path.join + + from src import InteractionDatasetSingleAgent, metrics def basestr(**kwargs): @@ -16,7 +19,7 @@ def basestr(**kwargs): """ return 'base_' -def main(method='bc', train=False, test=False, loc=0, **kwargs): +def main(method='bc', train=False, test=False, loc=0, config_path=None, **kwargs): """ Main loop for training and testing different imitation models Args: @@ -24,13 +27,26 @@ def main(method='bc', train=False, test=False, loc=0, **kwargs): 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 + config_file (str): path to config file + kwargs (dict): remaining kwargs for training loop """ + # get/set seed + seed = kwargs.get('seed',0) + torch.manual_seed(seed) + + # make prefix of output files outdir = opj('output',method,'loc%02i'%(loc)) if not os.path.isdir(outdir): os.mkdir(outdir) filestr = opj(outdir, basestr(**kwargs)) + # load config + if config_path: + with open(config_path, 'r') as cfg: + config = json5.load(cfg) + else: + raise Exception('No config path specified') + # method-based training if method=='bc': from src import bc @@ -41,14 +57,17 @@ def main(method='bc', train=False, test=False, loc=0, **kwargs): # default train / cv / test split datasets if train: + + # make policy, train and test datasets, and send to + policy = policy_class(config) 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) + train_fn(train_dataset, cv_dataset, policy, filestr, **kwargs) if test: # load policy - policy = policy_class.load_model(filestr=filestr, **kwargs) + policy = policy_class.load_model(config, filestr) policy.eval() # simulate policy @@ -77,6 +96,7 @@ def simulate_policy(policy, loc=0, track=0, filestr=''): env.render() done = False while not done: + # get action action = policy(ob) @@ -95,6 +115,8 @@ def parse_args(): test (bool): whether to run test loop method (str): the method to try for imitation loc (int): the location index of the roundabout + config (str): config path + seed (int): RNG seed """ import argparse parser = argparse.ArgumentParser(description='Save Expert Trajectories') @@ -105,7 +127,11 @@ def parse_args(): parser.add_argument("--test", help="test model", action="store_true") parser.add_argument("--method", help="modeling method", - choices=['bc', 'gail', 'advil'], default='bc') + choices=['bc', 'gail', 'advil'], default='bc') + parser.add_argument("--config", help="config file path", + default='config/networks.json5', type=str) + parser.add_argument('--seed', default=0, type=int, + help='seed') parser.add_argument() parser.add_argument() args = parser.parse_args() @@ -113,7 +139,9 @@ def parse_args(): 'train'=args.train, 'test'=args.test, 'method'=args.method, - 'loc'=args.loc + 'loc'=args.loc, + 'config'=args.config, + 'seed'=args.seed } return kwargs