From a37995694d98a8ecc61330211cd36bb9c502c30b Mon Sep 17 00:00:00 2001
From: ebuehrle <43623224+ebuehrle@users.noreply.github.com>
Date: Mon, 28 Feb 2022 12:23:02 +0100
Subject: [PATCH] Load policy config in evaluation

---
 eval_experiments.py | 21 +++++++++++++++++++--
 shail-experiment.py |  7 +++++++
 src/eval_main.py    | 14 +++++++-------
 3 files changed, 33 insertions(+), 9 deletions(-)

diff --git a/eval_experiments.py b/eval_experiments.py
index f9f1c8e..53e3684 100644
--- a/eval_experiments.py
+++ b/eval_experiments.py
@@ -1,17 +1,23 @@
 import os
 from src.eval_main import eval_main
 from src.evaluation.utils import load_and_average
+import torch
+import json
+
+activations = [torch.nn.Tanh, torch.nn.LeakyReLU]
 
 def main(method:str='expert', folder:str=None, locations=[(0,0)], skip_running=False):
     
+    exclude_keys_from_policy_kwargs = {'learning_rate', 'learning_rate_decay', 'clip_ratio', 'iterations_per_epoch', 'option'}
     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 = 'NormalizedSafeOptionsEvalEnv'
         env_kwargs={'stop_on_collision':True, 'max_episode_steps':1000, 'safe_actions_collision_method': None, 'abort_unsafe_collision_method': None}
     elif method in ['shail']:
         env = 'NormalizedSafeOptionsEvalEnv'
@@ -23,7 +29,18 @@ def main(method:str='expert', folder:str=None, locations=[(0,0)], skip_running=F
 
     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))
+        files = [f for f in files if f.endswith('.pt')]
+        with open(os.path.join(folder, 'config.json'), 'rb') as f:
+            config = json.load(f)
+        print('%i policy files found in %s folder' %(len(files), folder))
+        print('found policy config', config['policy'])
+
+    policy_config = {k: v for k, v in config['policy'].items() if k not in exclude_keys_from_policy_kwargs}
+    policy_config['activation'] = activations[policy_config['activation']]
+    print('final policy config', policy_config)
+
+    policy_kwargs.update(policy_config)
+    print('final policy kwargs', policy_kwargs)
 
     if not skip_running:
         for policy_file in files:
diff --git a/shail-experiment.py b/shail-experiment.py
index 4195759..598efa9 100644
--- a/shail-experiment.py
+++ b/shail-experiment.py
@@ -247,6 +247,13 @@ if __name__ == '__main__':
             os.makedirs(savepath)
         
         import shutil
+
+        # save config
+        shutil.copyfile(
+            args.test,
+            os.path.join(savepath, 'config.json')
+        )
+
         for i in range(args.test_seeds):
             s = analysis._checkpoints[i]['config']['seed']
             check_dir = analysis._checkpoints[i]['logdir']
diff --git a/src/eval_main.py b/src/eval_main.py
index 2e50286..37e78bd 100644
--- a/src/eval_main.py
+++ b/src/eval_main.py
@@ -41,33 +41,33 @@ def load_policy(method:str,
     if method == 'idm':
         policy = IDMRulePolicy(env, **policy_kwargs)
     elif method == 'bc':
-        policy = SetPolicy(env.action_space.shape[-1])
+        policy = SetPolicy(env.action_space.shape[-1], **policy_kwargs)
         policy.load_state_dict(torch.load(policy_file, map_location=ml))
         policy.eval()
     elif method == 'gail-trpo':
-        policy = SetPolicy(env.action_space.shape[-1])
+        policy = SetPolicy(env.action_space.shape[-1], **policy_kwargs)
         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':
-        policy = SetPolicy(env.action_space.shape[-1])
+        policy = SetPolicy(env.action_space.shape[-1], **policy_kwargs)
         policy.load_state_dict(torch.load(policy_file, map_location=ml))
         policy.eval()
     elif method == 'rail':
         raise NotImplementedError
     elif method == 'hail-trpo':
-        policy = SetDiscretePolicy(env.action_space.n)
+        policy = SetMaskedDiscretePolicy(env.action_space.n, **policy_kwargs)
         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 == 'hail':
-        policy = SetDiscretePolicy(env.action_space.n)
+        policy = SetMaskedDiscretePolicy(env.action_space.n, **policy_kwargs)
         policy.load_state_dict(torch.load(policy_file, map_location=ml))
         policy.eval()
     elif method == 'shail-trpo':
-        policy = SetMaskedDiscretePolicy(env.action_space.n)
+        policy = SetMaskedDiscretePolicy(env.action_space.n, **policy_kwargs)
         policy(
             torch.zeros(env.observation_space['observation'].shape),
             torch.zeros(env.observation_space['safe_actions'].shape)
@@ -76,7 +76,7 @@ def load_policy(method:str,
         policy.load_state_dict(torch.load(policy_file, map_location=ml))
         policy.eval()
     elif method == 'shail':
-        policy = SetMaskedDiscretePolicy(env.action_space.n)
+        policy = SetMaskedDiscretePolicy(env.action_space.n, **policy_kwargs)
         policy.load_state_dict(torch.load(policy_file, map_location=ml))
         policy.eval()
     else: