Merge branch 'main' of github.com:sisl/InteractionImitation

scratch/etienne/intersimple/gail_options_image.py

@@ -16,6 +16,7 @@ import tempfile
 import pathlib
 from imitation.util import logger
 from stable_baselines3.common.env_util import make_vec_env
+from tqdm import tqdm
 
 model_name = 'gail_options_image'
 env_settings = {'agent': 51, 'width': 36, 'height': 36, 'm_per_px': 2}
@@ -33,17 +34,17 @@ class OptionsCnnPolicy(stable_baselines3.common.policies.ActorCriticCnnPolicy):
         values = self.value_net(latent_vf)
         return values, distribution.distribution
 
-    def predict(self, obs):
+    def predict(self, obs, eps=1e-6):
         s, m = obs['obs'], obs['mask']
         values, prior = self._prior_distribution(s)
-        posterior = Categorical(prior.probs * m)
+        posterior = Categorical((prior.probs + eps) * m)
         ch = posterior.sample()
         return ch, values, posterior.log_prob(ch)
 
-    def evaluate_actions(self, obs, ch):
+    def evaluate_actions(self, obs, ch, eps=1e-6):
         s, m = obs['obs'], obs['mask']
         values, prior = self._prior_distribution(s)
-        posterior = Categorical(prior.probs * m)
+        posterior = Categorical((prior.probs + eps) * m)
         return values, posterior.log_prob(ch), posterior.entropy() # additional values used by PPO.train
 
 class OptionsEnv(gym.Wrapper):
@@ -72,10 +73,10 @@ class OptionsEnv(gym.Wrapper):
             self.episode_start = False
             if self.done:
                 self.s = self.env.reset()
-                self.m = available_actions(self.env)
                 self.done = False
                 self.episode_start = True
 
+            self.m = available_actions(self.env)
             self.ch, self.value, self.log_prob = generator.policy.predict({
                 'obs': torch.tensor(self.s).unsqueeze(0).to(generator.policy.device),
                 'mask': torch.tensor(self.m).unsqueeze(0).to(generator.policy.device),
@@ -86,18 +87,16 @@ class OptionsEnv(gym.Wrapper):
 
             assert not self.done
             assert self.plan
-            assert feasible(self.env, self.plan, self.ch)
+            #assert feasible(self.env, self.plan, self.ch)
 
             while not self.done and self.plan and feasible(self.env, self.plan, self.ch):
                 self.a, self.plan = self.plan[0], self.plan[1:]
                 self.a = self.env._normalize(self.a)
                 self.nexts, _, self.done, _ = self.env.step(self.a)
-                self.nextm = available_actions(self.env)
 
                 self._after_step()
 
                 self.s = self.nexts
-                self.m = self.nextm
             
             yield from self._transitions()
 
@@ -132,6 +131,7 @@ class HLOptions(OptionsEnv):
         super().__init__(*args, **kwargs)
 
     def _after_choice(self):
+        self.obs = {'obs': np.copy(self.s), 'mask': np.copy(self.m)}
         self.r = 0
         self.steps = 0
 
@@ -146,7 +146,7 @@ class HLOptions(OptionsEnv):
     
     def _transitions(self):
         yield {
-            'obs': {'obs': self.s, 'mask': self.m},
+            'obs': self.obs,
             'action': self.ch,
             'reward': self.r.detach(),
             'episode_start': self.episode_start,
@@ -250,7 +250,7 @@ def check_future_collisions_exact(env, actions):
         feasible (torch.Tensor): tensor of shape (B,) indicating whether the respective action profiles are collision-free
     """
     # First check with simple circle collision check
-    states, collision_tensor = check_future_coqllisions_circle(env, actions)
+    states, collision_tensor = check_future_collisions_circle(env, actions)
     (B, T, nv, _) = states.shape
     # For those that have colliding circles, check exactly
     colliding_mask = ~collision_tensor
@@ -282,12 +282,15 @@ def check_future_collisions_exact(env, actions):
 
 def feasible(env, plan, ch):
     """Check if input profile is feasible given current `env` state. Action `ch=0` is safe fallback."""
-    
+    if ch == 0:
+        return True
+
     # zero pad plan - Take (T,) np plan and convert it to (T, nv, 1) torch.Tensor
     full_plan = torch.zeros(len(plan), env._env._nv, 1)
     full_plan[:, env._agent, 0] = torch.tensor(plan)
-    valid = check_future_collisions_fast(env, [full_plan]) # check_future_collisions_fast takes in B-list and outputs (B,) bool tensor
-    return ch == 0 or valid.item()
+    # valid = check_future_collisions_fast(env, [full_plan]) # check_future_collisions_fast takes in B-list and outputs (B,) bool tensor
+    valid = check_future_collisions_exact(env, [full_plan]) # check_future_collisions_fast takes in B-list and outputs (B,) bool tensor
+    return valid.item()
 
 def flatten_transitions(transitions):
     return {
@@ -323,7 +326,7 @@ def train_generator(env, generator, discriminator, num_samples):
     
     generator.train()
 
-def train(expert_data, epochs=10, expert_batch_size=32, generator_steps=2048, discount=0.99):
+def train(expert_data, epochs=20, expert_batch_size=32, generator_steps=1024, discount=0.99):
     env = NRasterized(**env_settings)
     env.discount = discount
 
@@ -356,7 +359,7 @@ def train(expert_data, epochs=10, expert_batch_size=32, generator_steps=2048, di
         generator.tensorboard_log,
     )
 
-    for _ in range(epochs):
+    for _ in tqdm(range(epochs)):
         train_discriminator(LLOptions(env), generator, discriminator, num_samples=expert_batch_size)
         train_generator(HLOptions(env), generator, discriminator, num_samples=generator_steps)
     
@@ -368,8 +371,6 @@ if __name__ == '__main__':
  
     with open("data/NormalizedIntersimpleExpertMu.001_NRasterizedAgent51w36h36mppx2.pkl", "rb") as f:
         trajectories = pickle.load(f)
-    import pdb 
-    pdb.set_trace()
     transitions = rollout.flatten_trajectories(trajectories)
     generator = train(transitions)
 
@@ -380,7 +381,7 @@ if __name__ == '__main__':
 
     env = RenderOptions(NRasterized(**env_settings))
 
-    for s in env.sample_ll(generator):
+    for s in env.sample_ll(model):
         if s['dones']:
             break