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

scratch/arec/intersimple/gail_options_image.py

@@ -18,7 +18,7 @@ from stable_baselines3.common.env_util import make_vec_env
 from tqdm import tqdm
 
 import logging 
-#logging.basicConfig(level=logging.DEBUG)
+logging.basicConfig(level=logging.DEBUG)
 
 ALL_OPTIONS = [(v,t) for v in [0,2,4,6,8] for t in [5, 10, 20]] # option 0 is safe fallback
 
@@ -216,13 +216,60 @@ def check_future_collisions_fast(env, actions):
 
     return (distance > min_distance).all(-1).all(-1)
 
+def check_future_collisions_circles(env, actions, n_circles:int=2):
+    """Checks whether `env._agent` would collide with other agents assuming `actions` as input.
+    
+    Vehicles are (over-)approximated by multiple circles.
+
+    Args:
+        env (gym.Env): current environment state
+        actions (list of torch.Tensor): list of B (T, nv, adims) T-length action profiles
+    Returns:
+        feasible (torch.Tensor): tensor of shape (B,) indicating whether the respective action profiles are collision-free
+    """
+    assert n_circles >= 2
+    B, (T, nv, _) = len(actions), actions[0].shape
+
+    states = torch.stack(env._env.propagate_action_profile(actions), axis=0)
+    assert states.shape == (B, T, nv, 5)
+    centers = states[:, :, :, :2]
+    psi = states[:, :, :, 3]
+    lon = torch.stack([psi.cos(), psi.sin()],dim=-1) # (B, T, nv, 2)
+
+    # offset between [-env._env.lengths+env._env.widths/2, env._env.lengths/2-env._env.widths/2]
+    back = (-env._env._lengths/2+env._env._widths/2).unsqueeze(-1) # (nv, 1)
+    length = (env._env._lengths-env._env._widths).unsqueeze(-1) # (nv, 1)
+    diff_d =  back + length*(torch.arange(n_circles)/(n_circles-1)).unsqueeze(0)  # (nv, n_circles) 
+    assert diff_d.shape == (nv, n_circles)
+    
+    offsets = diff_d[None, None, :, :, None] *  lon[:, :, :, None, :]
+    assert offsets.shape == (B, T, nv, n_circles, 2)
+
+    expanded_centers=centers.unsqueeze(-2) + offsets #(B, T, nv, n_circles, 2)
+    assert expanded_centers.shape == (B, T, nv, n_circles, 2)
+    agent_centers = expanded_centers[:,:,env._agent:env._agent+1,:,:] #(B, T, 1, n_circles, 2)
+    ds = expanded_centers.reshape((B, T, nv*n_circles, 1, 2)) - agent_centers #(B, T, nv*nc,1, 2) - (B, T, 1, nc, 2) = (B, T, nv*nc, nc, 2) 
+
+    distance = (ds**2).sum(-1).sqrt().reshape((B, T, nv, n_circles, n_circles)) # (B, T, nv, nc, nc)
+    distance = torch.where(distance.isnan(), np.inf*torch.ones_like(distance), distance) # only collide with spawned agents
+    distance[:, :, env._agent] = np.inf # cannot collide with itself
+    assert distance.shape == (B, T, nv, n_circles, n_circles)
+
+    radius = env._env._widths*np.sqrt(2) / 2
+    min_distance = radius[env._agent] + radius
+    min_distance = min_distance[None, None, :, None, None]
+    assert min_distance.shape == (1, 1, nv, 1, 1)
+
+    return (distance > min_distance).all(-1).all(-1).all(-1).all(-1)
+
 def feasible(env, plan, ch):
     """Check if input profile is feasible given current `env` state. Action `ch=0` is safe fallback."""
     
     # 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
+    # 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_circles(env, [full_plan])
     return ch == 0 or valid.item()
 
 def flatten_transitions(transitions):
@@ -314,7 +361,7 @@ if __name__ == '__main__':
         transitions,
         env_class=env_class,
         env_settings=env_settings,
-        epochs=10, 
+        epochs=2, 
         discrim_batch_size=32, 
         generator_steps=2048, 
         discount=0.99

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