huangfu/InteractionImitation
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Move files to src

Browse Source
This commit is contained in:
ebuehrle
2021-10-26 17:41:05 +02:00
parent 24b91d4eec
commit bcddf422f0
10 changed files with 35 additions and 343 deletions
Download Patch File Download Diff File Expand all files Collapse all files

154
scratch/etienne/intersimple/gail/collisions.py
View File

@@ -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)

101
scratch/etienne/intersimple/gail/discriminator.py
View File

@@ -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)

59
scratch/etienne/intersimple/gail/policy.py
View File

@@ -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

13
scratch/etienne/intersimple/gail_options_image.py
View File

@@ -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']:

44
scratch/etienne/intersimple/gail_options_image_random.py
View File

@@ -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
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)
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)
train(transitions)

0
scratch/etienne/intersimple/gail/test_discriminator.py → src/discriminator/test_discriminator.py
View File

2
scratch/etienne/intersimple/gail/options.py → src/gail/options.py
View File

@@ -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):

0
scratch/etienne/intersimple/gail/test_options.py → src/gail/test_options.py
View File

0
scratch/etienne/intersimple/gail/train.py → src/gail/train.py
View File

3
src/util/collisions.py
View File

@@ -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)