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

experiments/experiment.py

@@ -3,6 +3,16 @@ from functools import partial
 import os
 opj = os.path.join
 
+# set up ray tune
+import ray
+from ray import tune
+from ray.tune import Analysis, ExperimentAnalysis
+from ray.tune.schedulers import ASHAScheduler
+from hyperopt import hp
+from ray.tune.suggest.hyperopt import HyperOptSearch
+
+
+
 from src.main import basestr, main
 
 def parse_args():
@@ -72,18 +82,18 @@ def get_ray_config(method:str)->dict:
     """
     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]),
+            "lr": tune.loguniform(1e-5, 1e-3),
+            "weight_decay": tune.choice([0, 0.1]),
             "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([1,2,3]),
-            "head_hidden_dim": tune.choice([16,32,64]),
+            "deepsets_phi_hidden_n": tune.randint(1,5),
+            "deepsets_phi_hidden_dim": tune.lograndint(8,65),
+            "deepsets_latent_dim": tune.lograndint(8,129),
+            "deepsets_rho_hidden_n": tune.randint(0,3),
+            "deepsets_rho_hidden_dim": tune.lograndint(8,129),
+            "deepsets_output_dim": tune.lograndint(4,129),
+            "head_hidden_n": tune.randint(1,6),
+            "head_hidden_dim": tune.lograndint(16,257),
             "head_final_activation": tune.choice(['sigmoid', None]),
         }
     else:
@@ -108,10 +118,6 @@ if __name__ == '__main__':
         main(config, filestr=filestr, **kwargs)
 
     elif kwargs['ray'] and kwargs['train']:
-        # set up ray tune
-        import ray
-        from ray import tune
-        from ray.tune.schedulers import ASHAScheduler
 
         ray.shutdown() 
         ray.init(log_to_driver=False)
@@ -121,29 +127,31 @@ if __name__ == '__main__':
             main(full_config, filestr='exp', datadir=datadir, **kwargs)
         
         datadir = os.path.abspath('./expert_data')
+        
         ray_config = get_ray_config(kwargs['method'])
-        custom_scheduler = ASHAScheduler(
-            metric='cv_loss',
-            mode="min",
-            grace_period=25,
-        )
+        search = HyperOptSearch(ray_config, max_concurrent=8, metric='cv_loss',mode="min",)
+        custom_scheduler = ASHAScheduler(metric='cv_loss', mode="min", grace_period=15)
+
         analysis = tune.run(
             partial(ray_train, datadir=datadir),
-            config=ray_config,
+            #config=ray_config,
+            search_alg=search,
             scheduler=custom_scheduler,
             local_dir=outdir,
             #resources_per_trial={"cpu": 2},
-            time_budget_s=45*60,
-            num_samples=2,
+            time_budget_s=120*60,
+            num_samples=100,
         )
     elif kwargs['ray'] and kwargs['test']:
-        import ray
-        from ray.tune import Analysis, ExperimentAnalysis
         analysis = Analysis(outdir, default_metric="cv_loss", default_mode="min")
         config = analysis.get_best_config()
         filepath = analysis.get_best_logdir()
+        filestr = opj(filepath, 'exp')
+        config_path = filestr+'_config.json'
+        with open(config_path, 'r') as cfg:
+            config = json5.load(cfg)
         print("Best ray experiment:", filepath)
-        main(None, filestr=opj(filepath, 'exp'), **kwargs)
+        main(config, filestr=filestr, **kwargs)
     else:
         raise Exception('No valid config found')
 

requirements.txt

@@ -6,3 +6,4 @@ json5
 tqdm
 tensorboard
 ray[tune]
+hyperopt

src/expert_data.py

@@ -69,15 +69,18 @@ def generate_expert_data(path: str='expert_data', loc: int = 0, track:int = 0, *
     torch.save(actions, filestr+'_raw_actions.pt')
     process_expert_observations(obs, actions, filestr)
 
-def process_expert_observations(obs, actions, filestr, dtype=torch.float32):
+def process_expert_observations(obs, actions, filestr, remove_outliers=True, dtype=torch.float32):
     """
     Process the expert observations and save them as torch tensors
     Args:
         obs (list[dict]): lost of observations
         actions (torch.Tensor): (T, nv, a) tensor of actions
         filestr (str): base filename with which to save out observation tensors
+        remove_outliers (bool): whether to remove datapoints with acceleration above or below 5 m/s/s
+        dtype (torch.Type): type to convert data to
     """
-    data = {'state':[], 'action':[], 'relative_state':[], 'path_x':[], 'path_y':[]}
+    keys = ['state', 'action', 'relative_state', 'path_x', 'path_y']
+    data = {key:[] for key in keys}
     assert len(obs) == len(actions), 'non-matching action and observation lengths'
     T = len(obs)
     max_nv = 0
@@ -104,6 +107,11 @@ def process_expert_observations(obs, actions, filestr, dtype=torch.float32):
         data['relative_state'][i] = torch.cat((data['relative_state'][i], pad), dim=1)
     data['relative_state'] = torch.cat(data['relative_state']).type(dtype)
 
+    if remove_outliers:
+        non_outlier_indices = torch.nonzero(torch.abs(data['action'][:,0]) < 5)
+        for key in keys:
+            data[key] = data[key][non_outlier_indices[:,0]]    
+
     # mandate equal length
     assert len(data['state']) == len(data['relative_state']) \
         == len(data['action']) == len(data['path_x']) \

src/metrics.py

@@ -58,8 +58,6 @@ def visualize_distribution(true, pred, filestr):
     """
     nni1 = ~torch.isnan(true)
     nni2 = ~torch.isnan(pred)
-    import pdb
-    pdb.set_trace()
     plt.figure()
     plt.hist(true[nni1].numpy(), density=True, bins=20)
     plt.hist(pred[nni2].numpy(), density=True, bins=20)