removing gail-trpo since performance is about the same as gail, adding experiment evaluation script, updating metric averaging to work

eval_experiments.py

@@ -0,0 +1,74 @@
+import os
+from src.eval_main import eval_main
+from src.evaluation.utils import load_and_average
+
+def main(method:str='expert', folder:str=None, locations=[(0,0)], skip_running=False):
+    
+    policy_kwargs = {}
+    if method in ['expert', 'idm']:
+        env, env_kwargs ='NRasterizedRouteIncrementingAgent', {}
+    elif method in ['bc','gail']:
+        env='NormalizedContinuousEvalEnv' 
+        env_kwargs={stop_on_collision:True, max_episode_steps:1000}
+    elif method in ['hail']:
+        env = 'NormalizedOptionsEvalEnv'
+        env_kwargs={stop_on_collision:True, max_episode_steps:1000}
+    elif method in ['shail']:
+        env = 'NormalizedSafeOptionsEvalEnv'
+        env_kwargs={stop_on_collision:True, max_episode_steps:1000}
+    else:
+        raise NotImplementedError
+
+    files = ['']
+
+    if folder is not None:
+        files = [os.path.join(folder, f) for f in os.listdir(folder) if os.path.isfile(os.path.join(folder, f))]
+        print('%i folders found in %s folder' %(len(files), folder))
+
+    if not skip_running:
+        for policy_file in files:
+            # run metrics on that file
+            outbase = eval_main(locations=locations,
+                                method=method, 
+                                policy_file=policy_file, 
+                                policy_kwargs=policy_kwargs,
+                                env=env, 
+                                env_kwargs=env_kwargs)
+        outfolder = os.path.dirname(outbase)
+    else:
+        locstr = 'loc_'+'_'.join([f'r{ro}t{tr}' for (ro,tr) in locations])
+        outfolder = os.path.join('out',method,locstr)
+    
+    import pdb
+    pdb.set_trace()
+    # load metrics from save_path
+    average_metrics = load_and_average(outfolder)
+    if method in ['expert', 'idm']:
+        latex_print(average_metrics, light=True)
+    else:
+        latex_print(average_metrics)
+
+def latex_print(am, light=False):
+    """
+    print latex line
+
+    am (Dict[str,tuple]): dict mapping metric_name to (mean, std)
+    """
+
+    print('success rate, distance travelled, RWSE_10, |DeltaV|, AccelJSD')
+    if light:
+        print("%2.1f& %2.1f & --- & --- & "
+            "--- \\\\" %( 100*am['success rate'][0], am['mean travel distance'][0]))
+        return
+    
+    print("%2.1f \\scriptstyle\\pm %2.1f & %2.1f \\scriptstyle\\pm %2.1f & "
+        "%1.2f \\scriptstyle\\pm %1.2f & %2.1f \\scriptstyle\\pm %1.1f & "
+        "%0.3f \\scriptstyle\\pm %0.3f \\\\" %( 100*am['success rate'][0], 100*am['success rate'][1],
+               am['mean travel distance'][0] , am['mean travel distance'][1]  ,
+               am['rwse_10s'][0] , am['rwse_10s'][1]  ,
+               am['average absolute average velocity'][0] , am['average absolute average velocity'][1]  ,
+               am['acceleration distribution divergence'][0] , am['acceleration distribution divergence'][1]  ))
+
+if __name__=='__main__':
+    import fire 
+    fire.Fire(main)

evaluate_models.sh

@@ -1,79 +1,14 @@
-# eval_main inputs
+# can add --skip_running if you've run before
-# locations: List[Tuple[int,int]]= [(0,0)],
-# method: str='expert', 
-# policy_file: str='', 
-# policy_kwargs: dict={},
-# env: str='NRasterizedRouteIncrementingAgent', 
-# env_kwargs: dict={},
-# seed: int=0
 
-# expert
+python -m eval_experiments 
-python -m src.eval_main
+python -m eval_experiments --locations='[(0,4)]'
-
+python -m eval_experiments --method idm
-# idm
+python -m eval_experiments --method idm --locations='[(0,4)]'
-python -m src.eval_main --method=idm
+python -m eval_experiments --method bc --folder='test_policies/bc/expA'
-
+python -m eval_experiments --method bc --folder='test_policies/bc/expB' --locations='[(0,4)]'
-# behavior cloning
+python -m eval_experiments --method gail --folder='test_policies/gail/expA'
-python -m src.eval_main --method=bc --policy_file='checkpoints/bc-intersimple-setobs2.pt' --env='NormalizedContinuousEvalEnv' --env_kwargs='{stop_on_collision:True,max_episode_steps:1000}' --seed=0
+python -m eval_experiments --method gail --folder='test_policies/gail/expB'--locations='[(0,4)]'
-#python -m src.eval_main --method=bc --policy_file='checkpoints/bc-intersimple-setobs2.pt' --env='NormalizedContinuousEvalEnv' --env_kwargs='{stop_on_collision:True,max_episode_steps:1000}' --seed=1
+python -m eval_experiments --method hail
-#python -m src.eval_main --method=bc --policy_file='checkpoints/bc-intersimple-setobs2.pt' --env='NormalizedContinuousEvalEnv' --env_kwargs='{stop_on_collision:True,max_episode_steps:1000}' --seed=2
+python -m eval_experiments --method hail --locations='[(0,4)]'
-#python -m src.eval_main --method=bc --policy_file='checkpoints/bc-intersimple-setobs2.pt' --env='NormalizedContinuousEvalEnv' --env_kwargs='{stop_on_collision:True,max_episode_steps:1000}' --seed=3
+python -m eval_experiments --method shail
-#python -m src.eval_main --method=bc --policy_file='checkpoints/bc-intersimple-setobs2.pt' --env='NormalizedContinuousEvalEnv' --env_kwargs='{stop_on_collision:True,max_episode_steps:1000}' --seed=4
+python -m eval_experiments --method shail --locations='[(0,4)]'
-#python -m src.evaluation.utils load_and_average out/bc
-
-# GAIL
-python -m src.eval_main --method=gail --policy_file='checkpoints/gail-intersimple-setobs2-03-02-22.pt' --env='NormalizedContinuousEvalEnv' --env_kwargs='{stop_on_collision:True,max_episode_steps:1000}' --seed=0
-#python -m src.eval_main --method=gail --policy_file='checkpoints/gail-intersimple-setobs2-03-02-22.pt' --env='NormalizedContinuousEvalEnv' --env_kwargs='{stop_on_collision:True,max_episode_steps:1000}' --seed=1
-#python -m src.eval_main --method=gail --policy_file='checkpoints/gail-intersimple-setobs2-03-02-22.pt' --env='NormalizedContinuousEvalEnv' --env_kwargs='{stop_on_collision:True,max_episode_steps:1000}' --seed=2
-#python -m src.eval_main --method=gail --policy_file='checkpoints/gail-intersimple-setobs2-03-02-22.pt' --env='NormalizedContinuousEvalEnv' --env_kwargs='{stop_on_collision:True,max_episode_steps:1000}' --seed=3
-#python -m src.eval_main --method=gail --policy_file='checkpoints/gail-intersimple-setobs2-03-02-22.pt' --env='NormalizedContinuousEvalEnv' --env_kwargs='{stop_on_collision:True,max_episode_steps:1000}' --seed=4
-#python -m src.evaluation.utils load_and_average out/gail
-
-# options GAIL
-python -m src.eval_main --method=ogail --policy_file='checkpoints/gail-options-setobs2-Feb15_18-49-05.pt' --env='NormalizedOptionsEvalEnv' --env_kwargs='{stop_on_collision:True,max_episode_steps:1000}'
-
-# options GAIL-PPO
-python -m src.eval_main --method=ogail-ppo --policy_file='checkpoints/gail-ppo-options-setobs2-Feb15_22-05-38.pt' --env='NormalizedOptionsEvalEnv' --env_kwargs='{stop_on_collision:True,max_episode_steps:1000}'
-
-# SHAIL
-python -m src.eval_main --method=sgail --policy_file='checkpoints/sgail-options-setobs2-Feb21_13-30-45.pt' --env='NormalizedSafeOptionsEvalEnv' --env_kwargs='{stop_on_collision:True,max_episode_steps:1000}'
-
-# SHAIL-PPO
-python -m src.eval_main --method=sgail-ppo --policy_file='checkpoints/sgail-ppo-options-setobs2-17-02-2022.pt' --env='NormalizedSafeOptionsEvalEnv' --env_kwargs='{stop_on_collision:True,max_episode_steps:1000}'
-
-
-
-### Same checkpoint files applied out of distribution (e.g. to track 5)
-# expert
-python -m src.eval_main --locations='[(0,4)]'
-
-# idm
-python -m src.eval_main --method=idm --locations='[(0,4)]'
-
-# behavior cloning
-python -m src.eval_main --method=bc --policy_file='checkpoints/bc-intersimple-setobs2.pt' --env='NormalizedContinuousEvalEnv' --env_kwargs='{stop_on_collision:True,max_episode_steps:1000}' --seed=0 --locations='[(0,4)]'
-#python -m src.eval_main --method=bc --policy_file='checkpoints/bc-intersimple-setobs2.pt' --env='NormalizedContinuousEvalEnv' --env_kwargs='{stop_on_collision:True,max_episode_steps:1000}' --seed=1 --locations='[(0,4)]'
-#python -m src.eval_main --method=bc --policy_file='checkpoints/bc-intersimple-setobs2.pt' --env='NormalizedContinuousEvalEnv' --env_kwargs='{stop_on_collision:True,max_episode_steps:1000}' --seed=2 --locations='[(0,4)]'
-#python -m src.eval_main --method=bc --policy_file='checkpoints/bc-intersimple-setobs2.pt' --env='NormalizedContinuousEvalEnv' --env_kwargs='{stop_on_collision:True,max_episode_steps:1000}' --seed=3 --locations='[(0,4)]'
-#python -m src.eval_main --method=bc --policy_file='checkpoints/bc-intersimple-setobs2.pt' --env='NormalizedContinuousEvalEnv' --env_kwargs='{stop_on_collision:True,max_episode_steps:1000}' --seed=4 --locations='[(0,4)]'
-#python -m src.evaluation.utils load_and_average out/bc
-
-# GAIL
-python -m src.eval_main --method=gail --policy_file='checkpoints/gail-intersimple-setobs2-03-02-22.pt' --env='NormalizedContinuousEvalEnv' --env_kwargs='{stop_on_collision:True,max_episode_steps:1000}' --seed=0 --locations='[(0,4)]'
-#python -m src.eval_main --method=gail --policy_file='checkpoints/gail-intersimple-setobs2-03-02-22.pt' --env='NormalizedContinuousEvalEnv' --env_kwargs='{stop_on_collision:True,max_episode_steps:1000}' --seed=1 --locations='[(0,4)]'
-#python -m src.eval_main --method=gail --policy_file='checkpoints/gail-intersimple-setobs2-03-02-22.pt' --env='NormalizedContinuousEvalEnv' --env_kwargs='{stop_on_collision:True,max_episode_steps:1000}' --seed=2 --locations='[(0,4)]'
-#python -m src.eval_main --method=gail --policy_file='checkpoints/gail-intersimple-setobs2-03-02-22.pt' --env='NormalizedContinuousEvalEnv' --env_kwargs='{stop_on_collision:True,max_episode_steps:1000}' --seed=3 --locations='[(0,4)]'
-#python -m src.eval_main --method=gail --policy_file='checkpoints/gail-intersimple-setobs2-03-02-22.pt' --env='NormalizedContinuousEvalEnv' --env_kwargs='{stop_on_collision:True,max_episode_steps:1000}' --seed=4 --locations='[(0,4)]'
-#python -m src.evaluation.utils load_and_average out/gail
-
-# options GAIL
-python -m src.eval_main --method=ogail --policy_file='checkpoints/gail-options-setobs2-Feb15_18-49-05.pt' --env='NormalizedOptionsEvalEnv' --env_kwargs='{stop_on_collision:True,max_episode_steps:1000}' --locations='[(0,4)]'
-
-# options GAIL-PPO
-python -m src.eval_main --method=ogail-ppo --policy_file='checkpoints/gail-ppo-options-setobs2-Feb15_22-05-38.pt' --env='NormalizedOptionsEvalEnv' --env_kwargs='{stop_on_collision:True,max_episode_steps:1000}' --locations='[(0,4)]'
-
-# SHAIL
-python -m src.eval_main --method=sgail --policy_file='checkpoints/sgail-options-setobs2-Feb21_13-30-45.pt' --env='NormalizedSafeOptionsEvalEnv' --env_kwargs='{stop_on_collision:True,max_episode_steps:1000}' --locations='[(0,4)]'
-
-# SHAIL-PPO
-python -m src.eval_main --method=sgail-ppo --policy_file='checkpoints/sgail-ppo-options-setobs2-17-02-2022.pt' --env='NormalizedSafeOptionsEvalEnv' --env_kwargs='{stop_on_collision:True,max_episode_steps:1000}' --locations='[(0,4)]'

gail-trpo-experiment.py

@@ -1,235 +0,0 @@
-# %%
-import os
-
-import gym
-from src.core.gail import gail, Buffer
-from src.core.value import SetValue
-from src.core.policy import SetPolicy
-from src.core.discriminator import DeepsetDiscriminator
-import torch
-
-from intersim.envs import IntersimpleLidarFlatRandom
-from intersim.envs.intersimple import speed_reward
-import functools
-from src.util.wrappers import CollisionPenaltyWrapper, TransformObservation, Setobs
-import numpy as np
-from torch.utils.tensorboard import SummaryWriter
-from ray import tune
-from datetime import datetime
-import json
-
-DIR = os.path.dirname(os.path.abspath(__file__))
-activations = [torch.nn.Tanh, torch.nn.LeakyReLU]
-
-obs_min = np.array([
-    [-1000, -1000, 0, -np.pi, -1e-1, 0.],
-    [0, -np.pi, -20, -20, -np.pi, -1e-1],
-    [0, -np.pi, -20, -20, -np.pi, -1e-1],
-    [0, -np.pi, -20, -20, -np.pi, -1e-1],
-    [0, -np.pi, -20, -20, -np.pi, -1e-1],
-    [0, -np.pi, -20, -20, -np.pi, -1e-1],
-]).reshape(-1)
-
-obs_max = np.array([
-    [1000, 1000, 20, np.pi, 1e-1, 0.],
-    [50, np.pi, 20, 20, np.pi, 1e-1],
-    [50, np.pi, 20, 20, np.pi, 1e-1],
-    [50, np.pi, 20, 20, np.pi, 1e-1],
-    [50, np.pi, 20, 20, np.pi, 1e-1],
-    [50, np.pi, 20, 20, np.pi, 1e-1],
-]).reshape(-1)
-
-def training_function(config):
-    np.random.seed(config['seed'])
-    torch.manual_seed(config['seed'])
-
-    if config['experiment'] == 'A':
-        envs = [Setobs(TransformObservation(CollisionPenaltyWrapper(
-            IntersimpleLidarFlatRandom(
-                n_rays=5,
-                reward=functools.partial(
-                    speed_reward,
-                    collision_penalty=0
-                ),
-                check_collisions=True,
-                stop_on_collision=config['trainenv']['stop_on_collision'],
-            ), collision_distance=6, collision_penalty=100),
-            lambda obs: (obs - obs_min) / (obs_max - obs_min + 1e-10)
-        )) for _ in range(60)]
-
-    elif config['experiment'] == 'B': 
-        envs = sum([[Setobs(TransformObservation(CollisionPenaltyWrapper(
-            IntersimpleLidarFlatRandom(
-                n_rays=5,
-                reward=functools.partial(
-                    speed_reward,
-                    collision_penalty=0
-                ),
-                check_collisions=True,
-                stop_on_collision=config['trainenv']['stop_on_collision'],
-                track=track,
-            ), collision_distance=6, collision_penalty=100),
-            lambda obs: (obs - obs_min) / (obs_max - obs_min + 1e-10)
-        )) for _ in range(15)] for track in range(4)],[])
-    
-    else:
-        raise NotImplementedError
-
-    env_fn = lambda i: envs[i]
-
-    policy = SetPolicy(env_fn(0).action_space.shape[0], 
-        n_hidden_layers=config['policy']['n_hidden_layers'],
-        hidden_layer_size=config['policy']['hidden_layer_size'],
-        activation=activations[config['policy']['activation']] ) # config net architecture
-
-    value = SetValue() # config net architecture
-    v_opt = torch.optim.Adam(value.parameters(), lr=config['value']['learning_rate'], weight_decay=config['value']['weight_decay'])
-
-    discriminator = DeepsetDiscriminator(
-        n_hidden_layers_element=config['discriminator']['n_hidden_layers_element'],
-        n_hidden_layers_global=config['discriminator']['n_hidden_layers_global'],
-        hidden_layer_size=config['discriminator']['hidden_layer_size'],
-        activation=activations[config['discriminator']['activation']],
-    )
-    disc_opt = torch.optim.Adam(discriminator.parameters(), lr=config['discriminator']['learning_rate'], weight_decay=config['discriminator']['weight_decay'])
-
-    if config['experiment'] == 'A': 
-        expert_data = torch.load(os.path.join(DIR, 'intersimple-expert-data-setobs2-loc0-track0.pt'))
-    elif config['experiment'] == 'B': 
-        expert_data = [
-            torch.load(os.path.join(DIR, 'intersimple-expert-data-setobs2-loc0-track0.pt')),
-            torch.load(os.path.join(DIR, 'intersimple-expert-data-setobs2-loc0-track1.pt')),
-            torch.load(os.path.join(DIR, 'intersimple-expert-data-setobs2-loc0-track2.pt')),
-            torch.load(os.path.join(DIR, 'intersimple-expert-data-setobs2-loc0-track3.pt')),
-        ]
-        d0 = [d[0] for d in expert_data]
-        d1 = [d[1] for d in expert_data]
-        d2 = [d[2] for d in expert_data]
-        d3 = [d[3] for d in expert_data]
-        expert_data = (torch.cat(d0), torch.cat(d1), torch.cat(d2), torch.cat(d3))
-    
-    expert_data = Buffer(*expert_data)
-
-    def callback(info):
-        tune.report(gen_mean_reward_per_episode=info['gen/mean_reward_per_episode'],
-                    disc_mean_reward_per_episode=info['disc/mean_reward_per_episode'], 
-                    mean_episode_length=info['gen/mean_episode_length'],
-                    gen_collision_rate=info['gen/collision_rate'])
-        
-        # save model checkpoints
-        ep = info['epoch'] + 1
-        if (ep % 25 == 0):
-            torch.save(info['policy'].state_dict(), f'policy_epoch{ep}.pt')
-
-    value, policy = gail(
-        env_fn=env_fn,
-        expert_data=expert_data,
-        discriminator=discriminator,
-        disc_opt=disc_opt,
-        disc_iters=config['discriminator']['iterations_per_epoch'],
-        policy=policy,
-        value=value,
-        v_opt=v_opt,
-        v_iters=config['value']['iterations_per_epoch'],
-        epochs=config['train_epochs'],
-        rollout_episodes=60, 
-        rollout_steps=200,
-        gamma=0.99,
-        gae_lambda=0.9,
-        delta=config['policy']['delta'],
-        backtrack_coeff=0.8,
-        backtrack_iters=10,
-        logger=SummaryWriter(comment='gail-trpo-options-setobs2'),
-        callback=callback,
-    )
-
-    # save model
-    torch.save(policy.state_dict(), 'policy_final.pt')
-
-if __name__ == '__main__':
-    import argparse
-    parser = argparse.ArgumentParser()
-    parser.add_argument('--train', choices=['A', 'B'])
-    parser.add_argument('--epochs', type=int, default=200)
-    parser.add_argument('--test', type=str, help='path to config file to run final training on')
-    parser.add_argument('--test_seeds', type=int, default=5)
-    parser.add_argument('--test_cpus', type=int, help='number of cpus available to split test seed training over')
-    args = parser.parse_args()
-
-    assert (args.train is None) ^ (args.test is None), 'Must either train on an experiment or test with a config file'
-
-    # if no test config specified, train
-    if args.test is None:
-        print('Running Tuning for Experiment %s'%(args.train))
-        analysis = tune.run(
-            training_function,
-            config={
-                'experiment': args.train,
-                'trainenv': {
-                    'stop_on_collision': False, 
-                },
-                'policy': {
-                    'learning_rate': 3e-4, 
-                    'learning_rate_decay': 1.0,
-                    'delta': 0.01,
-                    'iterations_per_epoch': 100,
-                    'hidden_layer_size': tune.grid_search([20, 40]),
-                    'n_hidden_layers': tune.grid_search([2, 3]), 
-                    'activation':0,
-                },
-                'value': {
-                    'learning_rate': 1e-4,
-                    'weight_decay': 1e-3, 
-                    'iterations_per_epoch': 1000, 
-                },
-                'discriminator': {
-                    'learning_rate': 1e-3, 
-                    'weight_decay': 1e-4, 
-                    'iterations_per_epoch': 100, 
-                    'n_hidden_layers_element': tune.grid_search([3,4]),
-                    'n_hidden_layers_global': tune.grid_search([1,2]),
-                    'hidden_layer_size': 10, 
-                    'activation': 0,
-                },
-                'train_epochs': args.epochs,
-                'seed': 0,
-            }
-        )
-        best_config = analysis.get_best_config(metric='gen_collision_rate', mode='min')
-        print('Best config: ', best_config)
-
-        # safe best_config
-        if not os.path.isdir(os.path.join(DIR, 'best_configs')):
-            os.mkdir(os.path.join(DIR, 'best_configs'))
-        
-        # save gail
-        with open(os.path.join(DIR, 'best_configs',f'gail-trpo_exp{args.train}.json'), 'w', encoding='utf-8') as f:
-            json.dump(best_config, f, ensure_ascii=False, indent=4)
-
-    # if config file specified, rerun it with appropriate number of seeds
-    else:
-        with open(args.test, 'rb') as f:
-            config = json.load(f)
-
-        print(f'Retraining {args.test} with {args.test_seeds} seeds on experiment {config["experiment"]}')
-
-        # rerun with appropriate number of seeds
-        rpt = {'cpu': int(args.test_cpus/args.test_seeds)} if (args.test_cpus is not None) else None
-        config['seed'] = tune.grid_search(list(range(1,args.test_seeds+1)))
-        analysis = tune.run(training_function, config=config, resources_per_trial=rpt)
-
-        # move final policies to appropriate directory
-        split_ = os.path.basename(args.test).split('_')
-        model = split_[0]
-        exper = split_[-1].split('.')[0]
-        savepath = os.path.join('test_policies',model,exper)
-        
-        if not os.path.isdir(savepath):
-            os.makedirs(savepath)
-        
-        import shutil
-        for i in range(args.test_seeds):
-            s = analysis._checkpoints[i]['config']['seed']
-            check_dir = analysis._checkpoints[i]['logdir']
-            shutil.copyfile(os.path.join(check_dir,'policy_final.pt'), 
-                os.path.join(savepath, f'policy_seed{s}.pt'))

src/eval_main.py

@@ -15,7 +15,7 @@ from src.options import envs as options_envs2
 from src.safe_options.policy import SetMaskedDiscretePolicy
 from src.safe_options import options as options_envs3
 from src.util.wrappers import IntersimpleTimeLimit
-
+import os
 from typing import Optional, List, Dict, Tuple
 import torch
 import numpy as np
@@ -44,29 +44,29 @@ def load_policy(method:str,
         policy = SetPolicy(env.action_space.shape[-1])
         policy.load_state_dict(torch.load(policy_file, map_location=ml))
         policy.eval()
-    elif method == 'gail':
+    elif method == 'gail-trpo':
         policy = SetPolicy(env.action_space.shape[-1])
         policy(torch.zeros(env.observation_space.shape))
         policy = ReparamPolicy(policy)
         policy.load_state_dict(torch.load(policy_file, map_location=ml))
         policy.eval()
-    elif method == 'gail-ppo':
+    elif method == 'gail':
         policy = SetPolicy(env.action_space.shape[-1])
         policy.load_state_dict(torch.load(policy_file, map_location=ml))
         policy.eval()
     elif method == 'rail':
         raise NotImplementedError
-    elif method == 'ogail':
+    elif method == 'hail-trpo':
         policy = SetDiscretePolicy(env.action_space.n)
         policy(torch.zeros(env.observation_space.shape))
         policy = ReparamPolicy(policy)
         policy.load_state_dict(torch.load(policy_file, map_location=ml))
         policy.eval()
-    elif method == 'ogail-ppo':
+    elif method == 'hail':
         policy = SetDiscretePolicy(env.action_space.n)
         policy.load_state_dict(torch.load(policy_file, map_location=ml))
         policy.eval()
-    elif method == 'sgail':
+    elif method == 'shail-trpo':
         policy = SetMaskedDiscretePolicy(env.action_space.n)
         policy(
             torch.zeros(env.observation_space['observation'].shape),
@@ -75,7 +75,7 @@ def load_policy(method:str,
         policy = ReparamSafePolicy(policy)
         policy.load_state_dict(torch.load(policy_file, map_location=ml))
         policy.eval()
-    elif method == 'sgail-ppo':
+    elif method == 'shail':
         policy = SetMaskedDiscretePolicy(env.action_space.n)
         policy.load_state_dict(torch.load(policy_file, map_location=ml))
         policy.eval()
@@ -375,6 +375,9 @@ def eval_main(
         policy_file (str): path to saved policy
         env (str): environment class
         method (str): method (expert, bc, gail, rail, hgail, hrail)
+
+    Returns:
+        outbase (str): string to outbase 
     """
     print(f'#############################################################################')
     print(f'Evaluating {method} from file {policy_file} on {env} at locations {locations}')
@@ -383,9 +386,17 @@ def eval_main(
     # set seed
     np.random.seed(seed)
     torch.manual_seed(seed)
-    pfilename = policy_file.split('/')[-1].split('.')[0]
     locstr = 'loc_'+'_'.join([f'r{ro}t{tr}' for (ro,tr) in locations])
-    outbase = f'out/{method}/{locstr}/{pfilename}_seed{seed}'
+    if policy_file == '':
+        method_path = method
+        name_base = method
+    else:
+        path_items =  policy_file.split('/')
+        name_base = path_items[-1].split('.')[0]
+        method_path = ('/').join(path_items[1:-1])
+    outfolder = os.path.join('out',method_path,locstr)
+    filebase = name_base + f'_tseed{seed}'
+    outbase = os.path.join(outfolder,filebase)
 
     # load expert metrics
     expert_metrics = generate_expert_metrics(locations)
@@ -404,6 +415,8 @@ def eval_main(
         save_metrics(smetrics, outbase+'_summary.pkl')
         cmetrics = comparison_metrics(policy_metrics, expert_metrics, outbase=outbase)
         save_metrics(cmetrics, outbase+'_comparison.pkl')
+    
+    return outbase
 
 if __name__=='__main__':
     import fire

src/evaluation/utils.py

@@ -34,13 +34,20 @@ def load_metrics(filestr:str):
         metrics = pickle.load(f) 
     return metrics
 
-def average_metrics(metric_list:List[Dict[str,float]]):
+def average_metrics(metric_list:List[Dict[str,float]], verbose:bool=True) ->Dict[str, tuple]:
     """
     Average all the metrics in the list
     
     Args:
         metric_list (list of dicts): list of metric dicts which each map a string to a float
+        verbose (bool): whether to print avg metrics
+    Returns:
+        average_metrics (Dict[str, tuple])
     """
+    average_metrics = {}
+    if len(metric_list) == 0:
+        return average_metrics
+    
     keys = list(metric_list[0].keys())
     N = len(metric_list)
     master_dict = {key:[] for key in keys}
@@ -50,29 +57,39 @@ def average_metrics(metric_list:List[Dict[str,float]]):
         master_dict[key] = np.array(master_dict[key])
         mu = np.nanmean(master_dict[key])
         std2 = np.nanstd(master_dict[key])*2
-        print(f'{key}: {mu} \pm {std2}')
+        if verbose:
+            print(f'{key}: {mu} \pm {std2}')
+        average_metrics[key] = (mu, std2)
+    return average_metrics
 
 
-def load_and_average(path:str):
+def load_and_average(path:str, verbose:bool=True):
     """
     Load and average all metric files in a particular folder
 
     Args:
         path (str)
+        verbose (bool): whether to print avg metrics
+    Returns:
+        avg_metrics (Dict[str, tuple])
     """
     assert os.path.isdir(path)
 
     # summary metrics
     summary_files = [os.path.join(path,f) for f in os.listdir(path) if f.endswith('summary.pkl')]
-    print(*summary_files, sep='\n')
+    if verbose: 
+        print(*summary_files, sep='\n')
     all_summary_metrics = [load_metrics(f) for f in summary_files]
-    average_metrics(all_summary_metrics)
+    avg_metrics = average_metrics(all_summary_metrics, verbose=verbose)
 
     # comparison metrics
     comp_files = [os.path.join(path,f) for f in os.listdir(path) if f.endswith('comparison.pkl')]
-    print(*comp_files, sep='\n')
+    if verbose:
+        print(*comp_files, sep='\n') 
     all_comp_metrics = [load_metrics(f) for f in comp_files]
-    average_metrics(all_comp_metrics)
+    comp_avg = average_metrics(all_comp_metrics, verbose=verbose)
+    avg_metrics.update(comp_avg)
+    return avg_metrics
 
 if __name__=='__main__':
     import fire