Integrate options env and policy

src/core/policy.py

@@ -15,8 +15,13 @@ class BasePolicy(nn.Module):
     def sample(self, dist):
         return self.torch_dist(dist).sample()
     
-    def predict(self, states):
-        return self.sample(self.forward(states))
+    def predict(self, observations, state=None, episode_start=None, deterministic=True):
+        observations = torch.tensor(observations)
+        if deterministic:
+            actions = self.forward(observations)[..., :self.action_dim]
+        else:
+            actions = self.sample(self.forward(observations))
+        return actions, None
 
     def log_prob(self, dist, actions):
         return self.torch_dist(dist).log_prob(actions)
@@ -58,6 +63,14 @@ class DiscretePolicy(BasePolicy):
     
     def torch_dist(self, dist):
         return Categorical(logits=dist)
+    
+    def predict(self, observations, state=None, episode_start=None, deterministic=True):
+        observations = torch.tensor(observations)
+        if deterministic:
+            _, actions = self.forward(observations).max(-1)
+        else:
+            actions = self.sample(self.forward(observations))
+        return actions, None
 
 class SetPolicy(Policy):
 

src/eval_main.py

@@ -8,6 +8,9 @@ from src.baselines import IDMRulePolicy
 from src.evaluation import IntersimpleEvaluation
 import src.gail.options as options_envs
 from src.evaluation.metrics import divergence, visualize_distribution
+from src.core.policy import SetPolicy, SetDiscretePolicy
+from src.core.reparam_module import ReparamPolicy
+from src.gail2 import envs as options_envs2
 
 from typing import Optional, List, Dict, Tuple
 import torch
@@ -33,12 +36,31 @@ def load_policy(method:str,
     if method == 'idm':
         policy = IDMRulePolicy(env, **policy_kwargs)
     elif method == 'bc':
-        raise NotImplementedError
+        policy = SetPolicy(env.action_space.shape[-1])
+        policy.load_state_dict(torch.load(policy_file))
+        policy.eval()
     elif method == 'gail':
-        policy = sb3.PPO.load(policy_file)
-        raise NotImplementedError
+        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))
+        policy.eval()
+    elif method == 'gail-ppo':
+        policy = SetPolicy(env.action_space.shape[-1])
+        policy.load_state_dict(torch.load(policy_file))
+        policy.eval()
     elif method == 'rail':
         raise NotImplementedError
+    elif method == 'ogail':
+        policy = SetDiscretePolicy(env.action_space.n)
+        policy(torch.zeros(env.observation_space.shape))
+        policy = ReparamPolicy(policy)
+        policy.load_state_dict(torch.load(policy_file))
+        policy.eval()
+    elif method == 'ogail-ppo':
+        policy = SetDiscretePolicy(env.action_space.n)
+        policy.load_state_dict(torch.load(policy_file))
+        policy.eval()
     elif method == 'sgail':
         policy = sb3.PPO.load(policy_file)
         raise NotImplementedError
@@ -158,6 +180,7 @@ def evaluate_policy(locations:List[Tuple[int,int]],
     """
     envs_dict = dict(intersim.envs.intersimple.__dict__)
     envs_dict.update(dict(options_envs.__dict__))
+    envs_dict.update(dict(options_envs2.__dict__))
     policy_metrics = [None]* len(locations)
 
     # iterate through vehicles

src/evaluation/evaluation.py

@@ -6,6 +6,7 @@ from typing import Callable, Dict, Optional
 import os
 import pickle
 from tqdm import tqdm
+from src.gail2.envs import OptionsEnv
 
 class IntersimpleEvaluation:
     """
@@ -34,6 +35,7 @@ class IntersimpleEvaluation:
         self.env = eval_env
         self.n_episodes = eval_env.nv
         self.use_pbar = use_pbar
+        self.is_options_env = isinstance(self.env, OptionsEnv)
 
         # metrics present on every step of every episode
         self.metric_keys_all = ['v_all', 'a_all', 'col_all']
@@ -85,11 +87,14 @@ class IntersimpleEvaluation:
         if self.use_pbar:
             self.pbar = tqdm(total=self.n_episodes)
 
+        if self.is_options_env:
+            print('Evaluating an options environment')
+
         evaluate_policy(
             policy, 
             self.env,
             n_eval_episodes=self.n_episodes,
-            callback=self.evaluate_policy_callback, 
+            callback=self.evaluate_options_policy_callback if self.is_options_env else self.evaluate_policy_callback,
             return_episode_rewards=False
         )
         if self.use_pbar:
@@ -100,6 +105,13 @@ class IntersimpleEvaluation:
             self.save(filestr)
         return self._metrics
 
+    def evaluate_options_policy_callback(self, local_vars, global_vars):
+        infos = local_vars['info']['ll']['infos']
+        dones = local_vars['info']['ll']['env_done']
+        agents = [info['agent'] for info in infos]
+        for info, done, agent in zip(infos, dones, agents):
+            self.eval_policy_step(info, done, agent)
+
     def evaluate_policy_callback(self, local_vars, global_vars):
         """
         Callback run in evaluate_policy after taking an action and receiving an observation
@@ -112,6 +124,9 @@ class IntersimpleEvaluation:
         env = local_vars['env'].envs[venv_i]
         assert isinstance(env, Intersimple)
 
+        self.eval_policy_step(info, done, _agent)
+
+    def eval_policy_step(self, info, done, _agent):
         # Increase collision counter if episode terminated with a collision
         self._metrics['v_all'][_agent].append(info['prev_state'][_agent,2].item())
         self._metrics['a_all'][_agent].append(info['action_taken'][_agent,0].item())

src/gail2/envs.py

@@ -1,6 +1,6 @@
 import gym
 import numpy as np
-from wrappers import Setobs, TransformObservation
+from src.gail2.wrappers import Wrapper, Setobs, TransformObservation
 from intersim.envs import IntersimpleLidarFlatIncrementingAgent
 
 obs_min = np.array([
@@ -30,7 +30,7 @@ def NormalizedOptionsEvalEnv(**kwargs):
         ), lambda obs: (obs - obs_min) / (obs_max - obs_min + 1e-10))
     ), options=[(0, 5), (1, 5), (2, 5), (4, 5), (6, 5), (8, 5)])
 
-class OptionsEnv(gym.Wrapper):
+class OptionsEnv(Wrapper):
 
     def __init__(self, env, options):
         super().__init__(env)