Add flag to skip seeds

This commit is contained in:
ebuehrle
2022-03-05 07:12:34 +01:00
parent a576f0fb18
commit 3fa370eb8a
2 changed files with 18 additions and 13 deletions

View File

@@ -6,7 +6,7 @@ import json
activations = [torch.nn.Tanh, torch.nn.LeakyReLU]
def main(method:str='expert', folder:str=None, locations=[(0,0)], skip_running=False, save_videos:bool=False, videos_folder:str='videos'):
def main(method:str='expert', folder:str=None, locations=[(0,0)], skip_running=False, save_videos:bool=False, videos_folder:str='videos', first_seed_only:bool=False):
exclude_keys_from_policy_kwargs = {'learning_rate', 'learning_rate_decay', 'clip_ratio', 'iterations_per_epoch', 'option'}
policy_kwargs = {}
@@ -30,8 +30,13 @@ 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))]
files = [f for f in files if f.endswith('.pt')]
if first_seed_only:
files = files[:1]
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'])