From 82407d522205f2b5f3f3bdedbb0fcd7d68587996 Mon Sep 17 00:00:00 2001 From: ebuehrle <43623224+ebuehrle@users.noreply.github.com> Date: Tue, 26 Oct 2021 12:37:10 +0200 Subject: [PATCH] Refactor options GAIL training script --- .../etienne/intersimple/gail/collisions.py | 154 +++++++++ scratch/etienne/intersimple/gail/options.py | 148 +++++++++ scratch/etienne/intersimple/gail/policy.py | 59 ++++ .../etienne/intersimple/gail/test_options.py | 59 ++++ scratch/etienne/intersimple/gail/train.py | 36 +++ .../etienne/intersimple/gail_options_image.py | 300 +----------------- 6 files changed, 463 insertions(+), 293 deletions(-) create mode 100644 scratch/etienne/intersimple/gail/collisions.py create mode 100644 scratch/etienne/intersimple/gail/options.py create mode 100644 scratch/etienne/intersimple/gail/policy.py create mode 100644 scratch/etienne/intersimple/gail/test_options.py create mode 100644 scratch/etienne/intersimple/gail/train.py diff --git a/scratch/etienne/intersimple/gail/collisions.py b/scratch/etienne/intersimple/gail/collisions.py new file mode 100644 index 0000000..f73bf5d --- /dev/null +++ b/scratch/etienne/intersimple/gail/collisions.py @@ -0,0 +1,154 @@ +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/options.py b/scratch/etienne/intersimple/gail/options.py new file mode 100644 index 0000000..eb89810 --- /dev/null +++ b/scratch/etienne/intersimple/gail/options.py @@ -0,0 +1,148 @@ +import gym +import torch +from .collisions import feasible +import numpy as np + +class OptionsEnv(gym.Wrapper): + + def __init__(self, env, options=[(v,t) for v in [0,2,4,6,8] for t in [5]], *args, **kwargs): + """option 0 is treated as safe fallback""" + + super().__init__(env, *args, **kwargs) + self.options = options + num_hl_options = len(self.options) + self.action_space = gym.spaces.Discrete(num_hl_options) + self.observation_space = gym.spaces.Dict({ + 'obs': env.observation_space, + 'mask': gym.spaces.Box(low=0, high=1, shape=(num_hl_options,)), + }) + + def _after_choice(self): + pass + + def _after_step(self): + pass + + def _transitions(self): + raise NotImplementedError('Use `LLOptions` or `HLOptions` for sampling.') + + def sample(self, generator): + self.done = True + while True: + self.episode_start = False + if self.done: + self.s = self.env.reset() + self.done = False + self.episode_start = True + + self.m = available_actions(self.env, self.options) + 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), + }) + self.plan = list(map(float, generate_plan(self.env, self.ch, self.options))) + + self._after_choice() + + assert not self.done + assert self.plan + #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._after_step() + + self.s = self.nexts + + yield from self._transitions() + +class LLOptions(OptionsEnv): + """Sample low-level (state, action) tuples for discriminator training.""" + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self.observation_space = self.observation_space['obs'] + + def _after_choice(self): + self._transition_buffer = [] + + def _after_step(self): + self._transition_buffer.append({ + 'obs': self.s, + 'next_obs': self.nexts, + 'acts': np.array((self.a,)), + 'dones': np.array(self.done), + }) + + def _transitions(self): + yield from self._transition_buffer + + def sample_ll(self, policy): + return self.sample(policy) + +class HLOptions(OptionsEnv): + """Sample high-level (state, action, reward) tuples for generator training.""" + + def __init__(self, *args, **kwargs): + 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 + + def _after_step(self): + self.r += self.discount**self.steps * self.discriminator.discrim_net.reward_train( + state=torch.tensor(self.s).unsqueeze(0).to(self.discriminator.discrim_net.device()), + action=torch.tensor([[self.a]]).to(self.discriminator.discrim_net.device()), + next_state=torch.tensor(self.s).unsqueeze(0).to(self.discriminator.discrim_net.device()), # unused + done=torch.tensor(self.done).unsqueeze(0).to(self.discriminator.discrim_net.device()), # unused + ) + self.steps += 1 + + def _transitions(self): + yield { + 'obs': self.obs, + 'action': self.ch, + 'reward': self.r.detach(), + 'episode_start': self.episode_start, + 'value': self.value.detach(), + 'log_prob': self.log_prob.detach(), + 'done': self.done, + } + + def sample_hl(self, policy, discriminator): + self.discriminator = discriminator + return self.sample(policy) + +class RenderOptions(LLOptions): + + def _after_step(self): + super()._after_step() + self.env.render() + + def close(self, *args, **kwargs): + self.env.close(*args, **kwargs) + +def available_actions(env, options): + """Return mask of available actions given current `env` state.""" + valid = np.array([feasible(env, generate_plan(env, i, options), i) for i in range(len(options))]) + return valid + +def target_velocity_plan(current_v: float, target_v: float, t: int, dt: float): + """Smoothly target a velocity in a given number of steps""" + # for now, constant acceleration + a = (target_v - current_v) / (t * dt) + return a*np.ones((t,)) + +def generate_plan(env, i, options): + """Generate input profile for high-level action `i`.""" + assert i < len(options), "Invalid option index {i}" + target_v, t = options[i] + current_v = env._env.state[env._agent, 1].item() # extract from env + plan = target_velocity_plan(current_v, target_v, t, env._env._dt) + assert len(plan) == t, "incorrect plan length" + return plan diff --git a/scratch/etienne/intersimple/gail/policy.py b/scratch/etienne/intersimple/gail/policy.py new file mode 100644 index 0000000..9bf71c9 --- /dev/null +++ b/scratch/etienne/intersimple/gail/policy.py @@ -0,0 +1,59 @@ +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/test_options.py b/scratch/etienne/intersimple/gail/test_options.py new file mode 100644 index 0000000..a2fdb83 --- /dev/null +++ b/scratch/etienne/intersimple/gail/test_options.py @@ -0,0 +1,59 @@ +import pickle +import imitation.data.rollout as rollout +from options import LLOptions, OptionsEnv +from intersim.envs import NRasterized +import itertools +import stable_baselines3 +from policy import OptionsCnnPolicy +from train import flatten_transitions +import numpy as np + +def test_ll_expert_data(): + with open("data/NormalizedIntersimpleExpertMu.001_NRasterizedAgent51w36h36mppx2.pkl", "rb") as f: + expert_trajectories = pickle.load(f) + expert_transitions = rollout.flatten_trajectories(expert_trajectories) + + env = LLOptions(NRasterized(agent=51, width=36, height=36, m_per_px=2)) + + gen_transitions = list(itertools.islice(env.sample_ll( + policy=stable_baselines3.PPO( + OptionsCnnPolicy, + OptionsEnv(env), + verbose=1, + ) + ), 10)) + gen_transitions = flatten_transitions(gen_transitions) + + assert expert_transitions[:10].obs.shape == gen_transitions['obs'].shape + assert expert_transitions[:10].next_obs.shape == gen_transitions['next_obs'].shape + assert expert_transitions[:10].acts.shape == gen_transitions['acts'].shape + assert expert_transitions[:10].dones.shape == gen_transitions['dones'].shape + +def test_ll_states(): + env = NRasterized() + policy = stable_baselines3.PPO( + OptionsCnnPolicy, + OptionsEnv(env), + verbose=1, + ) + llenv = LLOptions(env) + transitions = list(itertools.islice(llenv.sample_ll(policy=policy), 100)) + + env2 = NRasterized() + s2 = env2.reset() + for i, t in enumerate(transitions): + assert i == 0 or np.array_equal(t['obs'], transitions[i-1]['next_obs']) + assert np.array_equal(t['obs'], s2) + assert t['acts'].shape == (1,) + + nexts2, _, done2, _ = env2.step(t['acts']) + assert np.array_equal(t['next_obs'], nexts2) + assert np.array_equal(t['dones'], done2) + + if done2: + break + + s2 = nexts2 + +def test_hl_transitions(): + pass diff --git a/scratch/etienne/intersimple/gail/train.py b/scratch/etienne/intersimple/gail/train.py new file mode 100644 index 0000000..31edd46 --- /dev/null +++ b/scratch/etienne/intersimple/gail/train.py @@ -0,0 +1,36 @@ +import numpy as np +import itertools + +def flatten_transitions(transitions): + return { + 'obs': np.stack(list(t['obs'] for t in transitions), axis=0), + 'next_obs': np.stack(list(t['next_obs'] for t in transitions), axis=0), + 'acts': np.stack(list(t['acts'] for t in transitions), axis=0), + 'dones': np.stack(list(t['dones'] for t in transitions), axis=0), + } + +def train_discriminator(env, generator, discriminator, num_samples): + transitions = list(itertools.islice(env.sample_ll(generator), num_samples)) + generator_samples = flatten_transitions(transitions) + discriminator.train_disc(gen_samples=generator_samples) + +def train_generator(env, generator, discriminator, num_samples): + generator_samples = list(itertools.islice(env.sample_hl(generator, discriminator), num_samples+1)) + + generator.rollout_buffer.reset() + for s in generator_samples[:-1]: + generator.rollout_buffer.add( + obs=s['obs'], + action=s['action'].cpu(), + reward=s['reward'].cpu(), + episode_start=s['episode_start'], + value=s['value'], + log_prob=s['log_prob'], + ) + + generator.rollout_buffer.compute_returns_and_advantage( + last_values=generator_samples[-1]['value'], + dones=generator_samples[-1]['done'], + ) + + generator.train() diff --git a/scratch/etienne/intersimple/gail_options_image.py b/scratch/etienne/intersimple/gail_options_image.py index 31273b2..a26f07d 100644 --- a/scratch/etienne/intersimple/gail_options_image.py +++ b/scratch/etienne/intersimple/gail_options_image.py @@ -16,248 +16,14 @@ 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 model_name = 'gail_options_image' env_settings = {'agent': 51, 'width': 36, 'height': 36, 'm_per_px': 2} -ALL_OPTIONS = [(v,t) for v in [0,2,4,6,8] for t in [5, 10, 20]] # option 0 is safe fallback - -class OptionsCnnPolicy(stable_baselines3.common.policies.ActorCriticCnnPolicy): - - def __init__(self, observation_space, *args, **kwargs): - super().__init__(observation_space['obs'], *args, **kwargs) - - def _prior_distribution(self, s): - 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): - 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): - 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 - -class OptionsEnv(gym.Wrapper): - - def __init__(self, env, *args, **kwargs): - super().__init__(env, *args, **kwargs) - num_hl_options = len(ALL_OPTIONS) - self.action_space = gym.spaces.Discrete(num_hl_options) - self.observation_space = gym.spaces.Dict({ - 'obs': env.observation_space, - 'mask': gym.spaces.Box(low=0, high=1, shape=(num_hl_options,)), - }) - - def _after_choice(self): - pass - - def _after_step(self): - pass - - def _transitions(self): - raise NotImplementedError('Use `LLOptions` or `HLOptions` for sampling.') - - def sample(self, generator): - self.done = True - while True: - self.episode_start = False - if self.done: - self.s = self.env.reset() - 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), - }) - self.plan = list(map(float, generate_plan(self.env, self.ch))) - - self._after_choice() - - assert not self.done - assert self.plan - #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._after_step() - - self.s = self.nexts - - yield from self._transitions() - -class LLOptions(OptionsEnv): - """Sample low-level (state, action) tuples for discriminator training.""" - - def __init__(self, *args, **kwargs): - super().__init__(*args, **kwargs) - self.observation_space = self.observation_space['obs'] - - def _after_choice(self): - self._transition_buffer = [] - - def _after_step(self): - self._transition_buffer.append({ - 'obs': self.s, - 'next_obs': self.nexts, - 'acts': np.array((self.a,)), - 'dones': np.array(self.done), - }) - - def _transitions(self): - yield from self._transition_buffer - - def sample_ll(self, policy): - return self.sample(policy) - -class HLOptions(OptionsEnv): - """Sample high-level (state, action, reward) tuples for generator training.""" - - def __init__(self, *args, **kwargs): - 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 - - def _after_step(self): - self.r += self.discount**self.steps * self.discriminator.discrim_net.reward_train( - state=torch.tensor(self.s).unsqueeze(0).to(self.discriminator.discrim_net.device()), - action=torch.tensor([[self.a]]).to(self.discriminator.discrim_net.device()), - next_state=torch.tensor(self.s).unsqueeze(0).to(self.discriminator.discrim_net.device()), # unused - done=torch.tensor(self.done).unsqueeze(0).to(self.discriminator.discrim_net.device()), # unused - ) - self.steps += 1 - - def _transitions(self): - yield { - 'obs': self.obs, - 'action': self.ch, - 'reward': self.r.detach(), - 'episode_start': self.episode_start, - 'value': self.value.detach(), - 'log_prob': self.log_prob.detach(), - 'done': self.done, - } - - def sample_hl(self, policy, discriminator): - self.discriminator = discriminator - return self.sample(policy) - -class RenderOptions(LLOptions): - - def _after_step(self): - super()._after_step() - self.env.render() - - def close(self, *args, **kwargs): - self.env.close(*args, **kwargs) - -def available_actions(env): - """Return mask of available actions given current `env` state.""" - valid = np.array([feasible(env, generate_plan(env, i), i) for i in range(len(ALL_OPTIONS))]) - return valid - -def target_velocity_plan(current_v: float, target_v: float, t: int, dt: float): - """Smoothly target a velocity in a given number of steps""" - # for now, constant acceleration - a = (target_v - current_v) / (t * dt) - return a*np.ones((t,)) - -def generate_plan(env, i): - """Generate input profile for high-level action `i`.""" - assert i < len(ALL_OPTIONS), "Invalid option index {i}" - target_v, t = ALL_OPTIONS[i] - current_v = env._env.state[env._agent, 1].item() # extract from env - plan = target_velocity_plan(current_v, target_v, t, env._env._dt) - assert len(plan) == t, "incorrect plan length" - return plan - -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 - """ - 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) - - return (distance > min_distance).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.""" - 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 valid.item() - -def flatten_transitions(transitions): - return { - 'obs': np.stack(list(t['obs'] for t in transitions), axis=0), - 'next_obs': np.stack(list(t['next_obs'] for t in transitions), axis=0), - 'acts': np.stack(list(t['acts'] for t in transitions), axis=0), - 'dones': np.stack(list(t['dones'] for t in transitions), axis=0), - } - -def train_discriminator(env, generator, discriminator, num_samples): - transitions = list(itertools.islice(env.sample_ll(generator), num_samples)) - generator_samples = flatten_transitions(transitions) - discriminator.train_disc(gen_samples=generator_samples) - -def train_generator(env, generator, discriminator, num_samples): - generator_samples = list(itertools.islice(env.sample_hl(generator, discriminator), num_samples+1)) - - generator.rollout_buffer.reset() - for s in generator_samples[:-1]: - generator.rollout_buffer.add( - obs=s['obs'], - action=s['action'].cpu(), - reward=s['reward'].cpu(), - episode_start=s['episode_start'], - value=s['value'], - log_prob=s['log_prob'], - ) - - generator.rollout_buffer.compute_returns_and_advantage( - last_values=generator_samples[-1]['value'], - dones=generator_samples[-1]['done'], - ) - - generator.train() +ALL_OPTIONS = [(v,t) for v in [0,2,4,6,8] for t in [5, 10]] # option 0 is safe fallback def train(expert_data, epochs=20, expert_batch_size=32, generator_steps=1024, discount=0.99): env = NRasterized(**env_settings) @@ -280,7 +46,7 @@ def train(expert_data, epochs=20, expert_batch_size=32, generator_steps=1024, di generator = stable_baselines3.PPO( OptionsCnnPolicy, - OptionsEnv(env), + OptionsEnv(env, options=ALL_OPTIONS), verbose=1, n_steps=generator_steps, ) @@ -293,8 +59,8 @@ def train(expert_data, epochs=20, expert_batch_size=32, generator_steps=1024, di ) 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) + 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) return generator @@ -319,55 +85,3 @@ if __name__ == '__main__': break env.close(filestr='render/'+model_name) - -# %% Tests - -def test_ll_expert_data(): - with open("data/NormalizedIntersimpleExpertMu.001_NRasterizedAgent51w36h36mppx2.pkl", "rb") as f: - expert_trajectories = pickle.load(f) - expert_transitions = rollout.flatten_trajectories(expert_trajectories) - - env = LLOptions(NRasterized(agent=51, width=36, height=36, m_per_px=2)) - - gen_transitions = list(itertools.islice(env.sample_ll( - policy=stable_baselines3.PPO( - OptionsCnnPolicy, - OptionsEnv(env), - verbose=1, - ) - ), 10)) - gen_transitions = flatten_transitions(gen_transitions) - - assert expert_transitions[:10].obs.shape == gen_transitions['obs'].shape - assert expert_transitions[:10].next_obs.shape == gen_transitions['next_obs'].shape - assert expert_transitions[:10].acts.shape == gen_transitions['acts'].shape - assert expert_transitions[:10].dones.shape == gen_transitions['dones'].shape - -def test_ll_states(): - env = NRasterized() - policy = stable_baselines3.PPO( - OptionsCnnPolicy, - OptionsEnv(env), - verbose=1, - ) - llenv = LLOptions(env) - transitions = list(itertools.islice(llenv.sample_ll(policy=policy), 100)) - - env2 = NRasterized() - s2 = env2.reset() - for i, t in enumerate(transitions): - assert i == 0 or np.array_equal(t['obs'], transitions[i-1]['next_obs']) - assert np.array_equal(t['obs'], s2) - assert t['acts'].shape == (1,) - - nexts2, _, done2, _ = env2.step(t['acts']) - assert np.array_equal(t['next_obs'], nexts2) - assert np.array_equal(t['dones'], done2) - - if done2: - break - - s2 = nexts2 - -def test_hl_transitions(): - pass