Add SHAIL-PPO

2022-02-18 06:54:52 +01:00
parent 1624e1a349
commit 84351e77f2
4 changed files with 8 additions and 1 deletions
--- a/src/eval_main.py
+++ b/src/eval_main.py
@@ -72,6 +72,10 @@ def load_policy(method:str,
        policy = ReparamPolicy(policy)
        policy.load_state_dict(torch.load(policy_file))
        policy.eval()
+    elif method == 'sgail-ppo':
+        policy = SetMaskedDiscretePolicy(env.action_space.n)
+        policy.load_state_dict(torch.load(policy_file))
+        policy.eval()
    else:
        raise NotImplementedError
    return policy