updating main testing function to use configs and seeds, finishing first pass at behavior cloning policy and training loop. not yet tested

src/bc/bc.py

@@ -1,18 +1,28 @@
 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):
@@ -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
-
+        if 'action' in self._transforms.keys():
-        pass
+            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:
 
+        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
+            optimizer.zero_grad()
             loss.backwards()
             optimizer.step()
 
+            epoch_loss += loss.item() / len(train_dataset)
         
-        # measure L2 on every cv epoch
+        # Write epoch loss
+        if i % 10 == 0:
+            print('Epoch: {}, Training Loss: {}'.format(i, epoch_loss))
+
+        # 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)

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')
@@ -106,6 +128,10 @@ def parse_args():
         action="store_true") 
     parser.add_argument("--method", help="modeling method",
         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