Move files to src

scratch/etienne/intersimple/gail/collisions.py

@@ -1,154 +0,0 @@
-import torch
-import numpy as np
-from intersim.collisions import state_to_polygon
-
-def feasible(env, plan, ch, method='exact'):
-    """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)
-
-    # check_future_collisions_fast takes in B-list and outputs (B,) bool tensor
-    if method=='circle':
-        valid = check_future_collisions_fast(env, [full_plan]) 
-    elif method=='ncircles':
-        valid = check_future_collisions_ncircles(env, [full_plan])
-    elif method=='exact':
-        valid = check_future_collisions_exact(env, [full_plan])
-    else:
-        raise NotImplementedError('Invalid collision-checking method')
-    return valid.item()
-
-def check_future_collisions_ncircles(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 check_future_collisions_circle(env, actions):
-    """Compute collision information for circular vehicle approximations
-
-    Args:
-        env (gym.Env): current environment state
-        actions (list of torch.Tensor): list of B (T, nv, adims) T-length action profiles
-    Returns:
-        states (torch.Tensor): tensor of shape (B, T, nv, 5) of future states based on the action profiles
-        collision_tensor (torch.Tensor): tensor of shape (B, T, nv) of bools indicating which plan collides with which vehicles in which time frame
-            false: colliding, true: not colliding
-    """
-    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)
-
-    distance = ((states[:, :, :, :2] - states[:, :, env._agent:env._agent+1, :2])**2).sum(-1).sqrt()
-    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)
-
-    radius = (env._env._lengths**2 + env._env._widths**2).sqrt() / 2
-    min_distance = radius[env._agent] + radius
-    min_distance = min_distance.unsqueeze(0).unsqueeze(0)
-    assert min_distance.shape == (1, 1, nv)
-
-    collision_tensor = distance > min_distance
-    assert collision_tensor.shape == (B, T, nv)
-    return states, collision_tensor
-
-def check_future_collisions_fast(env, actions):
-    """Checks whether `env._agent` would collide with other agents assuming `actions` as input.
-    
-    Vehicles are (over-)approximated by single 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
-    """
-    _, collision_tensor = check_future_collisions_circle(env, actions)
-    return collision_tensor.all(-1).all(-1)
-
-def check_future_collisions_exact(env, actions):
-    """
-        Checks whether `env._agent` would collide with other agents assuming `actions` as input.
-
-    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
-    """
-    # First check with simple circle collision check
-    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
-    
-    ego_states = states[:, :, env._agent:env._agent+1, :].expand(states.shape)
-    assert ego_states.shape == states.shape
-
-    # get dimensions
-    lengths = env._env._lengths.expand(states.shape[:3])
-    widths = env._env._widths.expand(states.shape[:3])
-    ego_lengths = lengths[:, :, env._agent:env._agent+1].expand(lengths.shape)
-    ego_widths = widths[:, :, env._agent:env._agent+1].expand(widths.shape)
-    assert lengths.shape == widths.shape == ego_lengths.shape == ego_widths.shape == (B, T, nv)
-    
-    # For every collision instance between ego and other vehicle, check whether rectangles intersect
-    exact_collisions = torch.zeros_like(collision_tensor[colliding_mask])
-    for i, (ego_state, ego_length, ego_width, other_state, other_length, other_width) in enumerate(zip(
-        ego_states[colliding_mask], ego_lengths[colliding_mask], ego_widths[colliding_mask],
-        states[colliding_mask], lengths[colliding_mask], widths[colliding_mask]
-    )):
-        assert ego_state.shape == other_state.shape == (5,)
-        assert ego_length.shape == ego_width.shape == other_length.shape == other_width.shape == ()
-        p_ego = state_to_polygon(ego_state, ego_length, ego_width)
-        p_other = state_to_polygon(other_state, other_length, other_width)
-        exact_collisions[i] = p_ego.intersects(p_other)
-
-    collision_tensor[colliding_mask] = ~exact_collisions
-    return collision_tensor.all(-1).all(-1) 

scratch/etienne/intersimple/gail/discriminator.py

@@ -1,101 +0,0 @@
-import torch
-
-# imitation.rewards.discrim_nets.DiscrimNetGAIL is composed of self.discriminator (nn.Module),
-# which gets called with inputs (state, action) when needed.
-
-class CnnDiscriminator(torch.nn.Module):
-    """ConvNet similar to stable_baselines3.common.policies.ActorCriticCnnPolicy."""
-
-    def __init__(self, env):
-        super().__init__()
-
-        obs_channels, _, _ = env.observation_space.shape
-        (action_size,) = env.action_space.shape
-        in_channels = obs_channels + action_size
-
-        self.cnn = torch.nn.Sequential(
-            torch.nn.Conv2d(in_channels, 32, kernel_size=(8, 8), stride=(4, 4)), # 5+1 -> 32
-            torch.nn.ReLU(),
-            torch.nn.Conv2d(32, 64, kernel_size=(4, 4), stride=(2, 2)), # 32 -> 64
-            torch.nn.ReLU(),
-            torch.nn.Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1)), # 64 -> 64
-            torch.nn.ReLU(),
-            torch.nn.Flatten(start_dim=1, end_dim=-1),
-            torch.nn.LazyLinear(512), # 28224 -> 512
-            torch.nn.ReLU(),
-            torch.nn.LazyLinear(1), # 512 -> 1
-        )
-    
-    @staticmethod
-    def _concatenate(state, action):
-        b, _, h, w = state.shape
-        _, a = action.shape
-        act = action.unsqueeze(-1).unsqueeze(-1).expand((b, a, h, w))
-        sa = torch.cat((state, act), -3)
-        return sa
-
-    def forward(self, state, action):
-        sa = self._concatenate(state, action)
-        assert sa.ndim == 4
-        return self.cnn(sa).squeeze(1)
-
-class CnnDiscriminatorFlatAction(torch.nn.Module):
-    """ConvNet similar to stable_baselines3.common.policies.ActorCriticCnnPolicy."""
-
-    def __init__(self, env):
-        super().__init__()
-
-        obs_channels, _, _ = env.observation_space.shape
-        (action_size,) = env.action_space.shape
-        in_channels = obs_channels
-
-        self.cnn = torch.nn.Sequential(
-            torch.nn.Conv2d(in_channels, 32, kernel_size=(8, 8), stride=(4, 4)), # in_channels -> 32
-            torch.nn.ReLU(),
-            torch.nn.Conv2d(32, 64, kernel_size=(4, 4), stride=(2, 2)), # 32 -> 64
-            torch.nn.ReLU(),
-            torch.nn.Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1)), # 64 -> 64
-            torch.nn.ReLU(),
-            torch.nn.Flatten(start_dim=1, end_dim=-1),
-            torch.nn.LazyLinear(128), # 28224 -> 128
-        )
-        self.decoder = torch.nn.Sequential(
-            torch.nn.LazyLinear(64), #128 + 2 -> 64
-            torch.nn.ReLU(),
-            torch.nn.LazyLinear(64), #64 -> 64
-            torch.nn.ReLU(),
-            torch.nn.LazyLinear(1) #64 -> 1
-        )
-    
-    @staticmethod
-    def _concatenate(state, action):
-        b, s= state.shape
-        b, a = action.shape
-        sa = torch.cat((state, action), -1)
-        return sa
-
-    def forward(self, state, action):
-        s = self.cnn(state.float())
-        sa = self._concatenate(s, action)
-        assert sa.ndim == 2
-        return self.decoder(sa).squeeze(1)
-
-class MlpDiscriminator(torch.nn.Module):
-    """MLP similar to stable_baselines3.common.policies.ActorCriticPolicy."""
-
-    def __init__(self, env=None):
-        super().__init__()
-        self.flatten = torch.nn.Flatten(start_dim=1, end_dim=-1)
-        self.mlp = torch.nn.Sequential(
-            torch.nn.LazyLinear(64), # 42 -> 64
-            torch.nn.Tanh(),
-            torch.nn.LazyLinear(64), # 64 -> 64
-            torch.nn.Tanh(),
-            torch.nn.LazyLinear(1),  # 64 -> 1
-        )
-    
-    def forward(self, state, action):
-        flat = self.flatten(state)
-        sa = torch.cat((action, flat), -1)
-        assert sa.ndim == 2
-        return self.mlp(sa).squeeze(1)

scratch/etienne/intersimple/gail/policy.py

@@ -1,59 +0,0 @@
-import stable_baselines3
-from torch.distributions import Categorical
-
-class OptionsCnnPolicy(stable_baselines3.common.policies.ActorCriticCnnPolicy):
-    """
-    Class for high-level options policy (generator)
-    """
-    def __init__(self, observation_space, *args, **kwargs):
-        super().__init__(observation_space['obs'], *args, **kwargs)
-
-    def _prior_distribution(self, s):
-        """
-        Return prior distribution over high-level options (before masking)
-        Args:
-            s (torch.tensor): observation
-        Returns:
-            values (torch.tensor): values from critic
-            dist (torch.distributions): prior distribution over actions
-        """
-        latent_pi, latent_vf, latent_sde = self._get_latent(s)
-        distribution = self._get_action_dist_from_latent(latent_pi, latent_sde)
-        values = self.value_net(latent_vf)
-        return values, distribution.distribution
-
-    def predict(self, obs, eps=1e-6):
-        """
-        Will mask invalid states before making action selections
-        Args:
-            obs: dict with keys:
-                obs (torch.tensor): (*,o)  true observations
-                mask (torch.tensor): (*,m) mask over valid actions
-        Returns:
-            ch (torch.tensor): (*,a) sampled actions
-            values (torch.tensor): (*,) predicted value at observation
-            log_probs (torch.tensor): (*,) log probabilities of selected actions
-        """
-        s, m = obs['obs'], obs['mask']
-        values, prior = self._prior_distribution(s)
-        posterior = Categorical((prior.probs + eps) * m)
-        ch = posterior.sample()
-        return ch, values, posterior.log_prob(ch)
-
-    def evaluate_actions(self, obs, ch, eps=1e-6):
-        """
-        Evaluate particular actions
-        Args:
-            obs: dict with keys:
-                obs (torch.tensor): (*,o) true observations
-                mask (torch.tensor): (*,m) masks over valid actions
-            ch (torch.tensor): (*,a) selected actions
-        Returns:
-            values (torch.tensor): (*,) predicted value at observation
-            log_probs (torch.tensor): (*,) log probabilities of selected actions
-            ent (torch.tensor): (*,) entropy of each distribution over actions
-        """
-        s, m = obs['obs'], obs['mask']
-        values, prior = self._prior_distribution(s)
-        posterior = Categorical((prior.probs + eps) * m)
-        return values, posterior.log_prob(ch), posterior.entropy() # additional values used by PPO.train

scratch/etienne/intersimple/gail_options_image.py

@@ -1,5 +1,8 @@
 # %%
-from gail.discriminator import CnnDiscriminator, CnnDiscriminatorFlatAction
+import sys
+sys.path.append('../../../')
+
+from src.discriminator import CnnDiscriminator, CnnDiscriminatorFlatAction
 from imitation.algorithms import adversarial
 import stable_baselines3
 import torch.utils.data
@@ -16,9 +19,9 @@ import pathlib
 from imitation.util import logger
 from stable_baselines3.common.env_util import make_vec_env
 from tqdm import tqdm
-from gail.policy import OptionsCnnPolicy
+from src.policies.options import OptionsCnnPolicy
-from gail.options import OptionsEnv, LLOptions, HLOptions, RenderOptions
+from src.gail.options import OptionsEnv, LLOptions, HLOptions, RenderOptions
-from gail.train import train_discriminator, train_generator
+from src.gail.train import train_discriminator, train_generator
 
 model_name = 'gail_options_image'
 env_settings = {'agent': 51, 'width': 36, 'height': 36, 'm_per_px': 2}
@@ -78,7 +81,7 @@ if __name__ == '__main__':
     # %%
     model = stable_baselines3.PPO.load(model_name)
 
-    env = RenderOptions(NRasterized(**env_settings))
+    env = RenderOptions(NRasterized(**env_settings), options=ALL_OPTIONS)
 
     for s in env.sample_ll(model):
         if s['dones']:

scratch/etienne/intersimple/gail_options_image_random.py

@@ -1,5 +1,8 @@
 # %%
-from gail.discriminator import CnnDiscriminatorFlatAction
+import sys
+sys.path.append('../../../')
+
+from src.discriminator import CnnDiscriminatorFlatAction
 from imitation.algorithms import adversarial
 import stable_baselines3
 import torch.utils.data
@@ -16,16 +19,16 @@ import pathlib
 from imitation.util import logger
 from stable_baselines3.common.env_util import make_vec_env
 from tqdm import tqdm
-from gail.policy import OptionsCnnPolicy
+from src.policies.options import OptionsCnnPolicy
-from gail.options import OptionsEnv, LLOptions, HLOptions, RenderOptions
+from src.gail.options import OptionsEnv, LLOptions, HLOptions, RenderOptions
-from gail.train import train_discriminator, train_generator
+from src.gail.train import train_discriminator, train_generator
 
 model_name = 'gail_options_image_random'
 env_settings = {'width': 70, 'height': 70, 'm_per_px': 1}
 
 ALL_OPTIONS = [(v,t) for v in [0,2,4,6,8] for t in [5, 10, 20]] # option 0 is safe fallback
 
-def train(expert_data, epochs=20, expert_batch_size=32, generator_steps=1024, discount=0.99):
+def train(expert_data, epochs=100, expert_batch_size=64, generator_steps=1024, discount=0.99):
     env = NRasterizedRouteRandomAgent(**env_settings)
     env.discount = discount
 
@@ -61,27 +64,28 @@ def train(expert_data, epochs=20, expert_batch_size=32, generator_steps=1024, di
     for _ in tqdm(range(epochs)):
         train_discriminator(LLOptions(env, options=ALL_OPTIONS), generator, discriminator, num_samples=expert_batch_size)
         train_generator(HLOptions(env, options=ALL_OPTIONS), generator, discriminator, num_samples=generator_steps)
+        generator.save(model_name)
     
     return generator
 
+def video(model_name, env):
+    model = stable_baselines3.PPO.load(model_name)
+    env = RenderOptions(env, options=ALL_OPTIONS)
+    for s in env.sample_ll(model):
+        if s['dones']:
+            break
+    env.close(filestr='render/'+model_name)
+
+def evaluate():
+    video(
+        model_name=model_name,
+        env=NRasterizedRouteRandomAgent(**env_settings)
+    )
+
 # %%
 if __name__ == '__main__':
-    # %%
 
     with open("data/NormalizedIntersimpleExpertMu.001N10000_NRasterizedRouteRandomAgentw70h70mppx1.pkl", "rb") as f:
         trajectories = pickle.load(f)
     transitions = rollout.flatten_trajectories(trajectories)
-    generator = train(transitions, epochs=100)
+    train(transitions)
-
-    generator.save(model_name)
-
-    # %%
-    model = stable_baselines3.PPO.load(model_name)
-
-    env = RenderOptions(NRasterizedRouteRandomAgent(**env_settings))
-
-    for s in env.sample_ll(model):
-        if s['dones']:
-            break
-
-    env.close(filestr='render/'+model_name)

src/gail/options.py

@@ -1,6 +1,6 @@
 import gym
 import torch
-from .collisions import feasible
+from src.util.collisions import feasible
 import numpy as np
 
 class OptionsEnv(gym.Wrapper):

src/util/collisions.py

@@ -15,7 +15,7 @@ def feasible(env, plan, ch, method='exact'):
     if method=='circle':
         valid = check_future_collisions_fast(env, [full_plan]) 
     elif method=='ncircles':
-        valid = check_future_collisions_ncircles(env, [fullplan])
+        valid = check_future_collisions_ncircles(env, [full_plan])
     elif method=='exact':
         valid = check_future_collisions_exact(env, [full_plan])
     else:
@@ -152,4 +152,3 @@ def check_future_collisions_exact(env, actions):
 
     collision_tensor[colliding_mask] = ~exact_collisions
     return collision_tensor.all(-1).all(-1) 
-