getting hyperparameter tunning with ray tune working. updating default network with optimization and general parameters.

config/networks.json5

@@ -1,5 +1,5 @@
 {
-    ego_state: {
+    ego_encoder: {
         input_dim: 5, // number of state vars
         hidden_n: 0,
         hidden_dim: 5,
@@ -30,5 +30,13 @@
         hidden_dim: 50,
         output_dim: 1, // number of outputs e.g. number of actions, or just one
         final_activation: 'sigmoid',
-    }
+    },
+    optim: {
+            optimizer: 'adam',
+            lr: 1e-3,
+            weight_decay: 0.1,
+    },
+    train_epochs: 200,
+    train_batch_size: 32,
+    loss: 'huber',
 }

requirements.txt

@@ -5,3 +5,4 @@ pytest
 json5
 tqdm
 tensorboard
+ray[tune]

src/bc/__init__.py

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

src/bc/bc.py

@@ -7,10 +7,12 @@ from torch.utils.tensorboard import SummaryWriter
 from src.policies import DeepSetsPolicy
 from src.util.transform import MinMaxScaler
 from tqdm import tqdm
+import json5
+from ray import tune
 
 def bc_config(ray_config):
     config = {
-        'ego_state': {'input_dim': 5, 'hidden_n': 0, 'output_dim': 0},
+        'ego_encoder': {'input_dim': 5, 'hidden_n': 0, 'hidden_dim':0, 'output_dim': 0},
         'deepsets': {
             'input_dim': 5,
             'phi': {
@@ -21,7 +23,7 @@ def bc_config(ray_config):
             'rho': {'hidden_n': 0, 'hidden_dim': 10},
             'output_dim': 0
         },
-        'path_encoder': {'input_dim': 40, 'hidden_n': 0, 'output_dim': 0},
+        'path_encoder': {'input_dim': 40, 'hidden_n': 0, 'hidden_dim': 0, 'output_dim': 0},
         'head': {
             'input_dim': 0, # computed in constructor
             'hidden_n': ray_config['head_hidden_n'],
@@ -34,8 +36,8 @@ def bc_config(ray_config):
             'lr':ray_config['lr'],
             'weight_decay':ray_config['weight_decay']
         },
-        'train_epochs':1000,
+        'train_epochs': 100,
-        'train_batch_size': ray_config['batch_size'],
+        'train_batch_size': ray_config['train_batch_size'],
         'loss': ray_config['loss'],
 
     }
@@ -99,7 +101,7 @@ class BehaviorCloningPolicy():
         return action
 
     @classmethod
-    def load_model(cls, config: dict, filestr: str):
+    def load_model(cls, filestr: str, config: dict = None):
         """
         Load a model from a file prefix
         Args:
@@ -108,6 +110,9 @@ class BehaviorCloningPolicy():
         Returns
             model (BehaviorCloningPolicy): loaded model
         """
+        if not config:
+            with open(filestr+'_config.json', 'r') as cfg:
+                config = json5.load(cfg)
         transforms = pickle.load(open(filestr+'_transforms.pkl', 'rb'))
         model = cls(config, transforms=transforms)
         model._policy.load_state_dict(torch.load(filestr+'_model.pt'))
@@ -119,13 +124,17 @@ class BehaviorCloningPolicy():
     def parameters(self):
         return self._policy.parameters()
 
-    def save_model(self, filestr, save_transforms=True):
+    def save_model(self, filestr, save_config=True, save_transforms=True):
         """
         Save transforms and state_dict to a location specificed by filestr
         Args:
             filestr (str): string prefix to save model to
-            save_transforms (bool): whether to save transforms
+            save_config (bool): whether to save the config file (as a json)
+            save_transforms (bool): whether to save transforms (as a pickle)
         """
+        if save_config:
+            with open(filestr+'_config.json', 'w') as cfg:
+                json5.dump(self._config, cfg)
         if save_transforms:
             pickle.dump(self._transforms, open(filestr+'_transforms.pkl', 'wb'))
         torch.save(self._policy.state_dict(), filestr+'_model.pt')
@@ -148,15 +157,19 @@ def generate_transforms(dataset):
 
     return transforms
 
-def train(train_dataset, cv_dataset, policy, filestr, **kwargs):
+def train(config, policy, train_dataset, cv_dataset, filestr, **kwargs):
+    
+    using_ray = kwargs.get('ray', False)
 
     # hyperparams
-    train_epochs = 1000
+    loss_type = config['loss']
-    train_batch_size = 64
+    train_epochs = config['train_epochs']
-    learning_rate = 1e-3
+    train_batch_size = config['train_batch_size']
-    weight_decay = 0.1
+    optimizer_type = config['optim']['optimizer']
+    learning_rate = config['optim']['lr']
+    weight_decay = config['optim']['weight_decay']
 
-    cv_every = 10
+    cv_every = 5
     print_epoch_every = 1000
     print_cv_every = 1000
     checkpoint_every = 100
@@ -176,22 +189,30 @@ def train(train_dataset, cv_dataset, policy, filestr, **kwargs):
     policy.policy = policy.policy.type(train_dataset[0]['state'].dtype)
 
     # generate loss function, optimizer
+    cv_loss_fn = nn.MSELoss(reduction='sum')
+    if loss_type == 'huber':
         loss_fn = nn.HuberLoss(reduction='sum') 
+    elif loss_type == 'mse':
+        loss_fn = nn.MSELoss(reduction='sum')
+    else:
+        raise NotImplementedError
+    if optimizer_type == 'adam':  
         optimizer = torch.optim.Adam(policy.parameters(), lr=learning_rate, weight_decay=weight_decay)
+    else:
+        raise NotImplementedError
 
     # generate tensorboard writer
+    if not using_ray:
         writer = SummaryWriter(filestr)
     
     for i in tqdm(range(train_epochs)):
         
+        # train
         epoch_loss = 0
         for (batch_idx, batch) in enumerate(training_loader):
             
             # sample mini-batch and run through policy
             pred_action = policy(batch) 
-            if i == 0 and batch_idx==0:
-                pickle.dump(batch, open(filestr+'_test_batch.pkl', 'wb'))
-                policy.save_model(filestr)
             loss = loss_fn(pred_action, batch['action'])
 
             # compute loss and step optimizer
@@ -202,24 +223,34 @@ def train(train_dataset, cv_dataset, policy, filestr, **kwargs):
             epoch_loss += loss.item() / len(train_dataset)
         
         # Write epoch loss
-        
+        if using_ray:
+            tune.report(training_loss=epoch_loss, training_iteration=i)
+        else:
             writer.add_scalar('training loss',epoch_loss, i)
-        if i % print_epoch_every == 0:
+        
-            print('Epoch: {}, Training Loss: {}'.format(i, epoch_loss))
+        # if i % print_epoch_every == 0:
+        #     print('Epoch: {}, Training Loss: {}'.format(i, epoch_loss))
 
                 # measure cv loss
+
         if i % cv_every == 0:
             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'])
+                    loss = cv_loss_fn(pred_action, batch['action'])
                     cv_loss += loss.item() / len(cv_dataset)
+            
+            if using_ray:
+                tune.report(cv_loss=cv_loss, cv_epoch=i)
+            else:
                 writer.add_scalar('cv loss', cv_loss, i)
-        if i % print_cv_every == 0:    
+
-            print('Epoch: {}, CV Loss: {}'.format(i, cv_loss))
+        # if i % print_cv_every == 0:    
+        #     print('Epoch: {}, CV Loss: {}'.format(i, cv_loss))
 
         # save model checkpoints
         if i % checkpoint_every == 0:
             policy.save_model(filestr + '_epoch%04i'%(i) )
+
     policy.save_model(filestr)

src/main.py

@@ -6,6 +6,7 @@ import json5
 import os
 opj = os.path.join
 from tqdm import tqdm
+from functools import partial
 
 from src import InteractionDatasetSingleAgent, metrics
 from intersim.utils import get_map_path, get_svt
@@ -20,10 +21,11 @@ def basestr(**kwargs):
     """
     return 'base'
 
-def main(method='bc', train=False, test=False, loc=0, config_path=None, **kwargs):
+def main(config, method='bc', train=False, test=False, loc=0, datadir='./expert_data', filestr='', **kwargs):
     """
     Main loop for training and testing different imitation models
     Args:
+        config (dict): configuration dictionary for model
         train (bool): whether to run train loop
         test (bool): whether to run test loop
         method (str): the method to try for imitation
@@ -35,19 +37,6 @@ def main(method='bc', train=False, test=False, loc=0, config_path=None, **kwargs
     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.makedirs(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
@@ -61,9 +50,9 @@ def main(method='bc', train=False, test=False, loc=0, config_path=None, **kwargs
 
         # make policy, train and test datasets, and send to 
         policy = policy_class(config)
-        train_dataset = InteractionDatasetSingleAgent(loc=loc, tracks=[0,1,2])
+        train_dataset = InteractionDatasetSingleAgent(output_dir=datadir, loc=loc, tracks=[0])#,1,2])
-        cv_dataset = InteractionDatasetSingleAgent(loc=loc, tracks=[3])
+        cv_dataset = InteractionDatasetSingleAgent(output_dir=datadir, loc=loc, tracks=[3])
-        train_fn(train_dataset, cv_dataset, policy, filestr, **kwargs)
+        train_fn(config, policy, train_dataset, cv_dataset, filestr, **kwargs)
     
     if test:
 
@@ -76,7 +65,7 @@ def main(method='bc', train=False, test=False, loc=0, config_path=None, **kwargs
         simulate_policy(policy, loc=loc, track=track, filestr=filestr, nframes=500)
 
         # run test metrics
-        test_dataset = InteractionDatasetSingleAgent(loc=loc, tracks=[track])
+        test_dataset = InteractionDatasetSingleAgent(output_dir=datadir, loc=loc, tracks=[track])
         metrics(filestr, test_dataset, policy)
 
 
@@ -90,8 +79,9 @@ def simulate_policy(policy, loc=0, track=0, filestr='', nframes=float('inf')):
         filestr (str): path prefix to save simulation to
     """
     # animate from environment
-    svt, svt_path = get_svt(base='InteractionSimulator', loc=loc, track=track)
+    basepath = os.path.abspath('./InteractionSimulator')
-    osm = get_map_path(base='InteractionSimulator', loc=loc)
+    svt, svt_path = get_svt(base=basepath, loc=loc, track=track)
+    osm = get_map_path(base=basepath, loc=loc)
     env = gym.make('intersim:intersim-v0', svt=svt, map_path=osm, 
         min_acc=-np.inf, max_acc=np.inf)
     # env = gym.make('intersim:intersim-v0', loc=loc, track=track, 
@@ -134,12 +124,14 @@ def parse_args():
         help='location (default 0)')
     parser.add_argument("--train", help="train model",
         action="store_true")
+    parser.add_argument("--all-runs", help="use ray tune to run multiple experiments",
+        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("--config", help="config file path",
-        default='config/networks.json5', type=str)
+        default=None, type=str)
     parser.add_argument('--seed', default=0, type=int,
         help='seed')    
     args = parser.parse_args()
@@ -149,11 +141,103 @@ def parse_args():
         'method':args.method,
         'loc':args.loc,
         'config_path':args.config,
-        'seed':args.seed
+        'seed':args.seed,
+        'all_runs':args.all_runs
     }
     return kwargs
 
+def main_wrapper(**kwargs):
+    if kwargs['all_runs']:
+        pass
+    else:
+        main(**kwargs)
+
+def get_full_config(ray_config:dict, method:str)->dict:
+    """
+    Get full model configuration from ray config and method string
+    Args:
+        ray_config (dict): ray config
+        method (str): method to get full configuration for
+    """
+    if method == 'bc':
+        from src.bc import bc_config
+        config = bc_config(ray_config)
+    else:
+        raise NotImplementedError
+    return config
+
+def get_ray_config(method:str)->dict:
+    """
+    Get configuration for ray based on method.
+    Args: 
+        method (str): method to get configuration for
+    Returns:
+        ray_config (dict): configuration for ray
+    """
+    if method == 'bc':
+        ray_config = {
+            "lr": tune.choice([1e-4, 1e-3, 1e-2, 1e-1]),
+            "weight_decay": tune.choice([0.001, 0.01, 0.1, 0.5, 0.9]),
+            "loss": tune.choice(['huber', 'mse']),
+            "train_batch_size": tune.choice([16,32,64]),
+            "deepsets_phi_hidden_n": tune.choice([1,2,3]),
+            "deepsets_phi_hidden_dim": tune.choice([16,32,64]),
+            "deepsets_latent_dim": tune.choice([16,32,64]),
+            "deepsets_rho_hidden_n": tune.choice([0,1,2]),
+            "deepsets_rho_hidden_dim": tune.choice([16,32,64]),
+            "deepsets_output_dim": tune.choice([8,16,32,64]),
+            "head_hidden_n": tune.choice([0,1,2]),
+            "head_hidden_dim": tune.choice([16,32,64]),
+            "head_final_activation": tune.choice(['sigmoid', None]),
+        }
+    else:
+        raise NotImplementedError
+    return ray_config
 
 if __name__ == '__main__':
     kwargs = parse_args()
-    main(**kwargs)
+    
+    # make prefix of output files
+    outdir = opj('output',kwargs['method'],'loc%02i'%(kwargs['loc']))
+
+    if kwargs['config_path']:
+        # load config
+        with open(kwargs['config_path'], 'r') as cfg:
+            config = json5.load(cfg)
+        if not os.path.isdir(outdir):
+            os.makedirs(outdir)   
+        filestr = opj(outdir, basestr(**kwargs)) 
+        main(config, filestr=filestr, **kwargs)
+
+    elif kwargs['all_runs'] and kwargs['train']:
+        
+        def ray_train(config, datadir=None):
+            full_config = get_full_config(config, kwargs['method'])
+            main(full_config, filestr='exp', datadir=datadir, ray=True, **kwargs)
+        
+        # set up ray tune
+        from ray import tune
+        from ray.tune.schedulers import ASHAScheduler
+        datadir = os.path.abspath('./expert_data')
+        ray_config = get_ray_config(kwargs['method'])
+        custom_scheduler = ASHAScheduler(
+            metric='cv_loss',
+            mode="min",
+            grace_period=25,
+        )
+        analysis = tune.run(
+            partial(ray_train, datadir=datadir),
+            config=ray_config,
+            scheduler=custom_scheduler,
+            local_dir=outdir,
+            resources_per_trial={"cpu": 2},
+            num_samples=20,
+        )
+    else:
+        raise Exception('No valid config found')
+
+    
+    
+    
+
+    

src/policies/policy.py

@@ -14,13 +14,13 @@ class DeepSetsPolicy(Policy, nn.Module):
             config (dict): dictionary for configuring the deep sets policy
         """
         super(DeepSetsPolicy, self).__init__()
-        ego_config = config['ego_state']
+        ego_config = config['ego_encoder']
         deepsets_config = config['deepsets']
         pathnet_config = config['path_encoder']
 
-        self.ego_net = Phi.from_config(ego_config) if ego_config else lambda x: x
+        self.ego_net = Phi.from_config(ego_config)
-        self.deepsets_net = DeepSetsModule.from_config(deepsets_config) if deepsets_config else lambda x: x
+        self.deepsets_net = DeepSetsModule.from_config(deepsets_config)
-        self.path_net = Phi.from_config(pathnet_config) if pathnet_config else lambda x: x
+        self.path_net = Phi.from_config(pathnet_config)
 
         cat_dim = self.ego_net.output_dim + self.deepsets_net.output_dim + self.path_net.output_dim
         # head has number of concatenated features as input