diff --git a/scratch/etienne/intersimple/gail_options_image.py b/scratch/etienne/intersimple/gail_options_image.py index 71d4741..31273b2 100644 --- a/scratch/etienne/intersimple/gail_options_image.py +++ b/scratch/etienne/intersimple/gail_options_image.py @@ -1,40 +1,54 @@ # %% -import sys -sys.path.append('../../../') -from src.discriminator import CnnDiscriminator, CnnDiscriminatorFlatAction -from src.policies import OptionsCnnPolicy -from src.util import feasible -from src.data import load_experts - +from gail.discriminator import CnnDiscriminator, CnnDiscriminatorFlatAction from imitation.algorithms import adversarial -from imitation.util import logger -import imitation.data.rollout as rollout - import stable_baselines3 -from stable_baselines3.common.env_util import make_vec_env - -import torch import torch.utils.data import numpy as np +from intersim.envs.intersimple import NRasterized import itertools +from torch.distributions import Categorical import gym +import torch import pickle +import imitation.data.rollout as rollout import tempfile import pathlib +from imitation.util import logger +from stable_baselines3.common.env_util import make_vec_env from tqdm import tqdm -from intersim.envs.intersimple import NRasterized, NRasterizedRandomAgent, NRasterizedIncrementingAgent +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): - """ - Wrap an intersimple environment with an options generator - """ + def __init__(self, env, *args, **kwargs): - """ - Initialize wrapped environment and set high-level action and observation spaces - """ super().__init__(env, *args, **kwargs) num_hl_options = len(ALL_OPTIONS) self.action_space = gym.spaces.Discrete(num_hl_options) @@ -53,13 +67,6 @@ class OptionsEnv(gym.Wrapper): raise NotImplementedError('Use `LLOptions` or `HLOptions` for sampling.') def sample(self, generator): - """ - yield transitions using a generator - Args: - generator (sb3.PPO) - Yields: - - """ self.done = True while True: self.episode_start = False @@ -96,23 +103,13 @@ class LLOptions(OptionsEnv): """Sample low-level (state, action) tuples for discriminator training.""" def __init__(self, *args, **kwargs): - """ - LLOption uses the true LL observations - """ super().__init__(*args, **kwargs) - # overwrite observation space to just output obs directly self.observation_space = self.observation_space['obs'] def _after_choice(self): - """ - After each option choice, initialize/reset the transition buffer - """ self._transition_buffer = [] def _after_step(self): - """ - After each ll action, append s, s', a, done to transition buffer - """ self._transition_buffer.append({ 'obs': self.s, 'next_obs': self.nexts, @@ -121,18 +118,9 @@ class LLOptions(OptionsEnv): }) def _transitions(self): - """ - Yield from the transition buffer - """ yield from self._transition_buffer def sample_ll(self, policy): - """ - Args: - policy - Returns: - gen: iterable which samples low-level transitions from the environment - """ return self.sample(policy) class HLOptions(OptionsEnv): @@ -142,17 +130,11 @@ class HLOptions(OptionsEnv): super().__init__(*args, **kwargs) def _after_choice(self): - """ - After an option selection, initialize total reward and number of steps - """ self.obs = {'obs': np.copy(self.s), 'mask': np.copy(self.m)} self.r = 0 self.steps = 0 def _after_step(self): - """ - After each low-level action, add the discounted discriminated reward score (given a discriminator) - """ 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()), @@ -162,18 +144,6 @@ class HLOptions(OptionsEnv): self.steps += 1 def _transitions(self): - """ - Yield a single dictionary per high-level selected action - Fields: - obs: high-level state and mask at selection - action: chosen high-level action - reward: accumulated option reward - episode_start: whether the action was chosen at the episode start - value: the value estimate from the starting state - log_prob: the log_prob of the selected action from the starting state - done: whether the episode has ended - - """ yield { 'obs': self.obs, 'action': self.ch, @@ -185,29 +155,16 @@ class HLOptions(OptionsEnv): } def sample_hl(self, policy, discriminator): - """ - Args: - policy - discriminator: function with which to score rewards - Returns: - gen: iterable which samples high-level transitions from the environment - """ self.discriminator = discriminator return self.sample(policy) class RenderOptions(LLOptions): def _after_step(self): - """ - Render the environment after each low-level step - """ super()._after_step() self.env.render() def close(self, *args, **kwargs): - """ - On 'close', close the environment - """ self.env.close(*args, **kwargs) def available_actions(env): @@ -222,14 +179,7 @@ def target_velocity_plan(current_v: float, target_v: float, t: int, dt: float): return a*np.ones((t,)) def generate_plan(env, i): - """Generate input profile for high-level action `i`. - - Args: - env (gym.Env): current environment state - i (int): high-level action `i` - Returns: - plan (np.array): length T array of acceleration values - """ + """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 @@ -237,6 +187,44 @@ def generate_plan(env, i): 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), @@ -271,20 +259,8 @@ def train_generator(env, generator, discriminator, num_samples): generator.train() -def train(expert_data, env_class=NRasterizedRandomAgent, env_settings={}, epochs=10, discrim_batch_size=32, generator_steps=2048, discount=0.99): - """ - Args: - expert_data: list of transitions - env_class: environment class - env_settings: environment settings - epochs: number of epochs to train for - discrim_batch_size: discriminator batch size - generator_steps: number of steps taken in generator - discount: discount factor - Returns: - generator (stable_baselines3.PPO): options policy - """ - env = env_class(**env_settings) +def train(expert_data, epochs=20, expert_batch_size=32, generator_steps=1024, discount=0.99): + env = NRasterized(**env_settings) env.discount = discount tempdir = tempfile.TemporaryDirectory(prefix="quickstart") @@ -292,10 +268,10 @@ def train(expert_data, env_class=NRasterizedRandomAgent, env_settings={}, epochs logger.configure(tempdir_path / "GAIL/") print(f"All Tensorboards and logging are being written inside {tempdir_path}/.") - venv = make_vec_env(env_class, n_envs=1, env_kwargs=env_settings) + venv = make_vec_env(NRasterized, n_envs=1, env_kwargs=env_settings) discriminator = adversarial.GAIL( expert_data=expert_data, - expert_batch_size=discrim_batch_size, + expert_batch_size=expert_batch_size, discrim_kwargs={'discrim_net': CnnDiscriminatorFlatAction(venv)}, #discrim_kwargs={'discrim_net': CnnDiscriminator(venv)}, venv=venv, # unused @@ -317,7 +293,7 @@ def train(expert_data, env_class=NRasterizedRandomAgent, env_settings={}, epochs ) for _ in tqdm(range(epochs)): - train_discriminator(LLOptions(env), generator, discriminator, num_samples=discrim_batch_size) + train_discriminator(LLOptions(env), generator, discriminator, num_samples=expert_batch_size) train_generator(HLOptions(env), generator, discriminator, num_samples=generator_steps) return generator @@ -325,31 +301,18 @@ def train(expert_data, env_class=NRasterizedRandomAgent, env_settings={}, epochs # %% if __name__ == '__main__': # %% - model_name = 'gail_options_image' - env_class = NRasterizedRandomAgent - env_settings = {'width': 36, 'height': 36, 'm_per_px': 2} - - #env_class = NRasterized - #env_settings = {'agent': 51, 'width': 36, 'height': 36, 'm_per_px': 2} - files = ['../../../expert_data/DR_USA_Roundabout_FT0/track%04i/expert.pkl'%(i) for i in range(5)] - transitions=load_experts(files) - - generator = train( - transitions, - env_class=env_class, - env_settings=env_settings, - epochs=10, - discrim_batch_size=32, - generator_steps=2048, - discount=0.99 - ) + + with open("data/NormalizedIntersimpleExpertMu.001_NRasterizedAgent51w36h36mppx2.pkl", "rb") as f: + trajectories = pickle.load(f) + transitions = rollout.flatten_trajectories(trajectories) + generator = train(transitions) generator.save(model_name) # %% model = stable_baselines3.PPO.load(model_name) - env = RenderOptions(NRasterizedRandomAgent(**env_args)) + env = RenderOptions(NRasterized(**env_settings)) for s in env.sample_ll(model): if s['dones']: diff --git a/scratch/etienne/intersimple/gail_options_image_alltracks.py b/scratch/etienne/intersimple/gail_options_image_alltracks.py new file mode 100644 index 0000000..71d4741 --- /dev/null +++ b/scratch/etienne/intersimple/gail_options_image_alltracks.py @@ -0,0 +1,410 @@ +# %% +import sys +sys.path.append('../../../') +from src.discriminator import CnnDiscriminator, CnnDiscriminatorFlatAction +from src.policies import OptionsCnnPolicy +from src.util import feasible +from src.data import load_experts + +from imitation.algorithms import adversarial +from imitation.util import logger +import imitation.data.rollout as rollout + +import stable_baselines3 +from stable_baselines3.common.env_util import make_vec_env + +import torch +import torch.utils.data +import numpy as np +import itertools +import gym +import pickle +import tempfile +import pathlib +from tqdm import tqdm + +from intersim.envs.intersimple import NRasterized, NRasterizedRandomAgent, NRasterizedIncrementingAgent + +ALL_OPTIONS = [(v,t) for v in [0,2,4,6,8] for t in [5, 10, 20]] # option 0 is safe fallback + +class OptionsEnv(gym.Wrapper): + """ + Wrap an intersimple environment with an options generator + """ + def __init__(self, env, *args, **kwargs): + """ + Initialize wrapped environment and set high-level action and observation spaces + """ + 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): + """ + yield transitions using a generator + Args: + generator (sb3.PPO) + Yields: + + """ + 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): + """ + LLOption uses the true LL observations + """ + super().__init__(*args, **kwargs) + # overwrite observation space to just output obs directly + self.observation_space = self.observation_space['obs'] + + def _after_choice(self): + """ + After each option choice, initialize/reset the transition buffer + """ + self._transition_buffer = [] + + def _after_step(self): + """ + After each ll action, append s, s', a, done to transition buffer + """ + 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 the transition buffer + """ + yield from self._transition_buffer + + def sample_ll(self, policy): + """ + Args: + policy + Returns: + gen: iterable which samples low-level transitions from the environment + """ + 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): + """ + After an option selection, initialize total reward and number of steps + """ + self.obs = {'obs': np.copy(self.s), 'mask': np.copy(self.m)} + self.r = 0 + self.steps = 0 + + def _after_step(self): + """ + After each low-level action, add the discounted discriminated reward score (given a discriminator) + """ + 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 a single dictionary per high-level selected action + Fields: + obs: high-level state and mask at selection + action: chosen high-level action + reward: accumulated option reward + episode_start: whether the action was chosen at the episode start + value: the value estimate from the starting state + log_prob: the log_prob of the selected action from the starting state + done: whether the episode has ended + + """ + 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): + """ + Args: + policy + discriminator: function with which to score rewards + Returns: + gen: iterable which samples high-level transitions from the environment + """ + self.discriminator = discriminator + return self.sample(policy) + +class RenderOptions(LLOptions): + + def _after_step(self): + """ + Render the environment after each low-level step + """ + super()._after_step() + self.env.render() + + def close(self, *args, **kwargs): + """ + On 'close', close the environment + """ + 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`. + + Args: + env (gym.Env): current environment state + i (int): high-level action `i` + Returns: + plan (np.array): length T array of acceleration values + """ + 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 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() + +def train(expert_data, env_class=NRasterizedRandomAgent, env_settings={}, epochs=10, discrim_batch_size=32, generator_steps=2048, discount=0.99): + """ + Args: + expert_data: list of transitions + env_class: environment class + env_settings: environment settings + epochs: number of epochs to train for + discrim_batch_size: discriminator batch size + generator_steps: number of steps taken in generator + discount: discount factor + Returns: + generator (stable_baselines3.PPO): options policy + """ + env = env_class(**env_settings) + env.discount = discount + + tempdir = tempfile.TemporaryDirectory(prefix="quickstart") + tempdir_path = pathlib.Path(tempdir.name) + logger.configure(tempdir_path / "GAIL/") + print(f"All Tensorboards and logging are being written inside {tempdir_path}/.") + + venv = make_vec_env(env_class, n_envs=1, env_kwargs=env_settings) + discriminator = adversarial.GAIL( + expert_data=expert_data, + expert_batch_size=discrim_batch_size, + discrim_kwargs={'discrim_net': CnnDiscriminatorFlatAction(venv)}, + #discrim_kwargs={'discrim_net': CnnDiscriminator(venv)}, + venv=venv, # unused + gen_algo=stable_baselines3.PPO("CnnPolicy", venv), # unused + ) + + generator = stable_baselines3.PPO( + OptionsCnnPolicy, + OptionsEnv(env), + verbose=1, + n_steps=generator_steps, + ) + + # PPO.train requires logger as set up in + # PPO._setup_learn (called by PPO.learn) + generator._logger = stable_baselines3.common.utils.configure_logger( + generator.verbose, + generator.tensorboard_log, + ) + + for _ in tqdm(range(epochs)): + train_discriminator(LLOptions(env), generator, discriminator, num_samples=discrim_batch_size) + train_generator(HLOptions(env), generator, discriminator, num_samples=generator_steps) + + return generator + +# %% +if __name__ == '__main__': + # %% + model_name = 'gail_options_image' + env_class = NRasterizedRandomAgent + env_settings = {'width': 36, 'height': 36, 'm_per_px': 2} + + #env_class = NRasterized + #env_settings = {'agent': 51, 'width': 36, 'height': 36, 'm_per_px': 2} + files = ['../../../expert_data/DR_USA_Roundabout_FT0/track%04i/expert.pkl'%(i) for i in range(5)] + transitions=load_experts(files) + + generator = train( + transitions, + env_class=env_class, + env_settings=env_settings, + epochs=10, + discrim_batch_size=32, + generator_steps=2048, + discount=0.99 + ) + + generator.save(model_name) + + # %% + model = stable_baselines3.PPO.load(model_name) + + env = RenderOptions(NRasterizedRandomAgent(**env_args)) + + for s in env.sample_ll(model): + if s['dones']: + 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