huangfu/InteractionImitation
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

adding stablebaselines, adding safe loading for nonCuda cluster

Browse Source
This commit is contained in:
Arec Jamgochian
2022-02-24 23:18:02 -08:00
parent 5a5d8a7aff
commit 7feea74eb8
2 changed files with 11 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
requirements.txt
View File

@@ -7,4 +7,5 @@ tqdm
ray[tune] ray[tune]
hyperopt hyperopt
psutil psutil
fire fire
stable_baselines3

17
src/eval_main.py
View File

@@ -37,21 +37,22 @@ def load_policy(method:str,
Returns: Returns:
policy (Optional[BaseAlgorithm]): the policy to evaluate policy (Optional[BaseAlgorithm]): the policy to evaluate
""" """
ml = torch.device('cpu') if not torch.cuda.is_available() else None
if method == 'idm': if method == 'idm':
policy = IDMRulePolicy(env, **policy_kwargs) policy = IDMRulePolicy(env, **policy_kwargs)
elif method == 'bc': elif method == 'bc':
policy = SetPolicy(env.action_space.shape[-1]) policy = SetPolicy(env.action_space.shape[-1])
policy.load_state_dict(torch.load(policy_file)) policy.load_state_dict(torch.load(policy_file, map_location=ml))
policy.eval() policy.eval()
elif method == 'gail': elif method == 'gail':
policy = SetPolicy(env.action_space.shape[-1]) policy = SetPolicy(env.action_space.shape[-1])
policy(torch.zeros(env.observation_space.shape)) policy(torch.zeros(env.observation_space.shape))
policy = ReparamPolicy(policy) policy = ReparamPolicy(policy)
policy.load_state_dict(torch.load(policy_file)) policy.load_state_dict(torch.load(policy_file, map_location=ml))
policy.eval() policy.eval()
elif method == 'gail-ppo': elif method == 'gail-ppo':
policy = SetPolicy(env.action_space.shape[-1]) policy = SetPolicy(env.action_space.shape[-1])
policy.load_state_dict(torch.load(policy_file)) policy.load_state_dict(torch.load(policy_file, map_location=ml))
policy.eval() policy.eval()
elif method == 'rail': elif method == 'rail':
raise NotImplementedError raise NotImplementedError
@@ -59,11 +60,11 @@ def load_policy(method:str,
policy = SetDiscretePolicy(env.action_space.n) policy = SetDiscretePolicy(env.action_space.n)
policy(torch.zeros(env.observation_space.shape)) policy(torch.zeros(env.observation_space.shape))
policy = ReparamPolicy(policy) policy = ReparamPolicy(policy)
policy.load_state_dict(torch.load(policy_file)) policy.load_state_dict(torch.load(policy_file, map_location=ml))
policy.eval() policy.eval()
elif method == 'ogail-ppo': elif method == 'ogail-ppo':
policy = SetDiscretePolicy(env.action_space.n) policy = SetDiscretePolicy(env.action_space.n)
policy.load_state_dict(torch.load(policy_file)) policy.load_state_dict(torch.load(policy_file, map_location=ml))
policy.eval() policy.eval()
elif method == 'sgail': elif method == 'sgail':
policy = SetMaskedDiscretePolicy(env.action_space.n) policy = SetMaskedDiscretePolicy(env.action_space.n)
@@ -72,11 +73,11 @@ def load_policy(method:str,
torch.zeros(env.observation_space['safe_actions'].shape) torch.zeros(env.observation_space['safe_actions'].shape)
) )
policy = ReparamSafePolicy(policy) policy = ReparamSafePolicy(policy)
policy.load_state_dict(torch.load(policy_file)) policy.load_state_dict(torch.load(policy_file, map_location=ml))
policy.eval() policy.eval()
elif method == 'sgail-ppo': elif method == 'sgail-ppo':
policy = SetMaskedDiscretePolicy(env.action_space.n) policy = SetMaskedDiscretePolicy(env.action_space.n)
policy.load_state_dict(torch.load(policy_file)) policy.load_state_dict(torch.load(policy_file, map_location=ml))
policy.eval() policy.eval()
else: else:
raise NotImplementedError raise NotImplementedError
@@ -406,4 +407,4 @@ def eval_main(
if __name__=='__main__': if __name__=='__main__':
import fire import fire
fire.Fire(eval_main) fire.Fire(eval_main)