From bcddf422f09e376156bcb2f8a45c35426c330cf0 Mon Sep 17 00:00:00 2001 From: ebuehrle <43623224+ebuehrle@users.noreply.github.com> Date: Tue, 26 Oct 2021 17:41:05 +0200 Subject: [PATCH] Move files to src --- .../etienne/intersimple/gail/collisions.py | 154 ------------------ .../etienne/intersimple/gail/discriminator.py | 101 ------------ scratch/etienne/intersimple/gail/policy.py | 59 ------- .../etienne/intersimple/gail_options_image.py | 13 +- .../intersimple/gail_options_image_random.py | 46 +++--- .../discriminator}/test_discriminator.py | 0 .../intersimple => src}/gail/options.py | 2 +- .../intersimple => src}/gail/test_options.py | 0 .../etienne/intersimple => src}/gail/train.py | 0 src/util/collisions.py | 3 +- 10 files changed, 35 insertions(+), 343 deletions(-) delete mode 100644 scratch/etienne/intersimple/gail/collisions.py delete mode 100644 scratch/etienne/intersimple/gail/discriminator.py delete mode 100644 scratch/etienne/intersimple/gail/policy.py rename {scratch/etienne/intersimple/gail => src/discriminator}/test_discriminator.py (100%) rename {scratch/etienne/intersimple => src}/gail/options.py (99%) rename {scratch/etienne/intersimple => src}/gail/test_options.py (100%) rename {scratch/etienne/intersimple => src}/gail/train.py (100%) diff --git a/scratch/etienne/intersimple/gail/collisions.py b/scratch/etienne/intersimple/gail/collisions.py deleted file mode 100644 index f73bf5d..0000000 --- a/scratch/etienne/intersimple/gail/collisions.py +++ /dev/null @@ -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) diff --git a/scratch/etienne/intersimple/gail/discriminator.py b/scratch/etienne/intersimple/gail/discriminator.py deleted file mode 100644 index 1c9664c..0000000 --- a/scratch/etienne/intersimple/gail/discriminator.py +++ /dev/null @@ -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) diff --git a/scratch/etienne/intersimple/gail/policy.py b/scratch/etienne/intersimple/gail/policy.py deleted file mode 100644 index 9bf71c9..0000000 --- a/scratch/etienne/intersimple/gail/policy.py +++ /dev/null @@ -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 diff --git a/scratch/etienne/intersimple/gail_options_image.py b/scratch/etienne/intersimple/gail_options_image.py index a26f07d..3dd83c8 100644 --- a/scratch/etienne/intersimple/gail_options_image.py +++ b/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 gail.options import OptionsEnv, LLOptions, HLOptions, RenderOptions -from gail.train import train_discriminator, train_generator +from src.policies.options import OptionsCnnPolicy +from src.gail.options import OptionsEnv, LLOptions, HLOptions, RenderOptions +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']: diff --git a/scratch/etienne/intersimple/gail_options_image_random.py b/scratch/etienne/intersimple/gail_options_image_random.py index 23d39cc..216ea9f 100644 --- a/scratch/etienne/intersimple/gail_options_image_random.py +++ b/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 gail.options import OptionsEnv, LLOptions, HLOptions, RenderOptions -from gail.train import train_discriminator, train_generator +from src.policies.options import OptionsCnnPolicy +from src.gail.options import OptionsEnv, LLOptions, HLOptions, RenderOptions +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 -# %% -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) - - generator.save(model_name) - - # %% +def video(model_name, env): model = stable_baselines3.PPO.load(model_name) - - env = RenderOptions(NRasterizedRouteRandomAgent(**env_settings)) - + 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) + train(transitions) diff --git a/scratch/etienne/intersimple/gail/test_discriminator.py b/src/discriminator/test_discriminator.py similarity index 100% rename from scratch/etienne/intersimple/gail/test_discriminator.py rename to src/discriminator/test_discriminator.py diff --git a/scratch/etienne/intersimple/gail/options.py b/src/gail/options.py similarity index 99% rename from scratch/etienne/intersimple/gail/options.py rename to src/gail/options.py index eb89810..21dd7af 100644 --- a/scratch/etienne/intersimple/gail/options.py +++ b/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): diff --git a/scratch/etienne/intersimple/gail/test_options.py b/src/gail/test_options.py similarity index 100% rename from scratch/etienne/intersimple/gail/test_options.py rename to src/gail/test_options.py diff --git a/scratch/etienne/intersimple/gail/train.py b/src/gail/train.py similarity index 100% rename from scratch/etienne/intersimple/gail/train.py rename to src/gail/train.py diff --git a/src/util/collisions.py b/src/util/collisions.py index 8327857..f73bf5d 100644 --- a/src/util/collisions.py +++ b/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) -