From e7b0aea4271ad3572f3d8a9cadaf3eb497f84295 Mon Sep 17 00:00:00 2001 From: ebuehrle <43623224+ebuehrle@users.noreply.github.com> Date: Fri, 10 Sep 2021 16:15:34 +0200 Subject: [PATCH] run options gail --- scratch/etienne/intersimple/gail_options.py | 136 ---------- .../etienne/intersimple/gail_options_image.py | 250 ++++++++++++++++++ 2 files changed, 250 insertions(+), 136 deletions(-) delete mode 100644 scratch/etienne/intersimple/gail_options.py create mode 100644 scratch/etienne/intersimple/gail_options_image.py diff --git a/scratch/etienne/intersimple/gail_options.py b/scratch/etienne/intersimple/gail_options.py deleted file mode 100644 index 7890a35..0000000 --- a/scratch/etienne/intersimple/gail_options.py +++ /dev/null @@ -1,136 +0,0 @@ -from gail.discriminator import MlpDiscriminator -from imitation.algorithms import adversarial -import stable_baselines3 -import torch.utils.data -import numpy as np -from intersim.envs.intersimple import Intersimple -import itertools -from torch.distributions import Categorical -import gym -import torch - -ALL_OPTIONS = [(v,t) for v in [0,2,4,6,8] for t in [5, 10, 20, 50, 100]] - -class OptionsMlpPolicy: - - def __init__(self, *args, **kwargs): - self._policy = stable_baselines3.common.policies.ActorCriticPolicy( - *args, **kwargs - ) - - def _prior_distribution(self, s): - latent_pi, _, latent_sde = self._policy._get_latent(s) - distribution = self._policy._get_action_dist_from_latent(latent_pi, latent_sde) - return distribution.distribution - - def predict(self, obs): - s, m = obs - prior = self._prior_distribution(s) - posterior = Categorical(prior.probs * m) - ch = posterior.sample() - return ch - - def evaluate_actions(self, obs, ch): - s, m = obs - values = self._policy.value_net(s) - prior = self._prior_distribution(s) - posterior = Categorical(prior.probs * m) - return values, posterior.logprob(ch), posterior.entropy() # additional values used by PPO.train - -def available_actions(env): - """Return mask of available actions given current `env` state.""" - valid = np.array([feasible(env, generate_plan(env, i)) for i in range(len(ALL_OPTIONS))]) - return valid - -def target_velocity_plan(current_v: float, target_v: float, t: int) - """Smoothly target a velocity in a given number of steps""" - # for now, constant acceleration - a = (target_v - current_v) / t - 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) - assert len(plan) == t, "incorrect plan length" - return plan - -def feasible(env, plan): - """Check if input profile is feasible given current `env` state.""" - - # 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 = env._env.check_future_collisions([full_plan]) # check_future_collisions takes in B-list and outputs (B,) bool tensor - return valid.item() - -def sample_ll(env, generator): - """Sample low-level (state, action) pairs for discriminator training.""" - done = True - while True: - if done: - s = env.reset() - - m = available_actions(env) - ch = generator.policy.predict((s, m)) - plan = list(generate_plan(env, ch)) - - while not done and plan and feasible(env, plan): - a = plan.pop() - yield (s, a) - s, _, done, _ = env.step(a) - -def train_discriminator(env, expert_data, generator, discriminator, generator_batch_size): - expert_samples = next(expert_data) - generator_samples = itertools.islice(sample_ll(env, generator), generator_batch_size) - discriminator.train_disc(expert_samples, generator_samples) - -def sample_hl(env, generator, discriminator): - """Sample high-level (state, action, reward) tuples for generator training.""" - done = True - while True: - if done: - s = env.reset() - - m = available_actions(env) - obs = (s, m) - ch = generator.policy.predict((s, m)) - plan = list(generate_plan(env, ch)) - r = 0 - discount = 1 - - while not done and plan and feasible(env, plan): - a = plan.pop() - r += discount * discriminator.discrim_net(s, a) - discount *= env.discount - s, _, done, _ = env.step(a) - - yield (obs, ch, r) - -def train_generator(env, generator, discriminator, generator_batch_size): - generator_samples = itertools.islice(sample_hl(env, generator, discriminator), generator_batch_size) - generator.rollout_buffer.reset() - generator.rollout_buffer.add(generator_samples) - generator.train() - -class OptionsEnv(gym.Wrapper): - - def __init__(self, env): - super().__init__(env) - self.action_space = gym.spaces.Discrete(env.num_hl_options) - -def train(expert_data, epochs=10, generator_batch_size=1024, expert_batch_size=1024, num_hl_options=10, num_hl_steps=10, discount=0.99): - env = Intersimple() - env.num_hl_options = num_hl_options - env.num_hl_steps = num_hl_steps - env.discount = discount - - discriminator = adversarial.GAIL(discrim_kwargs={'discrim_net': MlpDiscriminator()}) - generator = stable_baselines3.PPO(OptionsMlpPolicy, OptionsEnv(env)) - expert_data = torch.utils.data.DataLoader(expert_data, expert_batch_size) - - for _ in range(epochs): - train_discriminator(env, expert_data, generator, discriminator, generator_batch_size) - train_generator(env, generator, discriminator, generator_batch_size) diff --git a/scratch/etienne/intersimple/gail_options_image.py b/scratch/etienne/intersimple/gail_options_image.py new file mode 100644 index 0000000..4b3aaa5 --- /dev/null +++ b/scratch/etienne/intersimple/gail_options_image.py @@ -0,0 +1,250 @@ +from gail.discriminator import CnnDiscriminator +from imitation.algorithms import adversarial +import stable_baselines3 +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 + +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]] + +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): + s, m = obs['obs'], obs['mask'] + values, prior = self._prior_distribution(s) + posterior = Categorical(prior.probs * m) + ch = posterior.sample() + return ch, values, posterior.log_prob(ch) + + def evaluate_actions(self, obs, ch): + s, m = obs['obs'], obs['mask'] + values, prior = self._prior_distribution(s) + posterior = Categorical(prior.probs * m) + return values, posterior.log_prob(ch), posterior.entropy() # additional values used by PPO.train + +def available_actions(env): + """Return mask of available actions given current `env` state.""" + valid = np.array([feasible(env, generate_plan(env, 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 feasible(env, plan): + """Check if input profile is feasible given current `env` state.""" + + # 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 = env._env.check_future_collisions([full_plan]) # check_future_collisions takes in B-list and outputs (B,) bool tensor + return valid.item() + +def sample_ll(env, generator): + """Sample low-level (state, action) pairs for discriminator training.""" + done = True + while True: + if done: + s = env.reset() + done = False + + m = available_actions(env) + ch, _, _ = generator.policy.predict({'obs': torch.tensor(s).unsqueeze(0), 'mask': torch.tensor(m).unsqueeze(0)}) + plan = list(generate_plan(env, ch)) + + while not done and plan and feasible(env, plan): + a, plan = plan[0], plan[1:] + nexts, _, done, _ = env.step(a) + yield { + 'obs': s, + 'next_obs': nexts, + 'acts': np.array((a,)), + 'dones': np.array(done), + } + s = nexts + +def sample_hl(env, generator, discriminator): + """Sample high-level (state, action, reward) tuples for generator training.""" + done = True + while True: + episode_start = False + if done: + s = env.reset() + m = available_actions(env) + done = False + episode_start = True + + obs = {'obs': s, 'mask': m} + ch, value, log_prob = generator.policy.predict({'obs': torch.tensor(s).unsqueeze(0), 'mask': torch.tensor(m).unsqueeze(0)}) + plan = list(map(float, generate_plan(env, ch))) + r = 0 + steps = 0 + + while not done and plan and feasible(env, plan): + a, plan = plan[0], plan[1:] + r += discriminator.discrim_net.discriminator(torch.tensor(s).unsqueeze(0), torch.tensor([[a]])) + steps += 1 + s, _, done, _ = env.step(a) + m = available_actions(env) + + yield { + 'obs': obs, + 'action': ch, + 'reward': r.detach() / steps, + 'episode_start': episode_start, + 'value': value.detach(), + 'log_prob': log_prob.detach(), + 'done': done, + } + +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(sample_ll(env, 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(sample_hl(env, 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'], + reward=s['reward'], + 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() + +class OptionsEnv(gym.Wrapper): + + def __init__(self, env): + super().__init__(env) + 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 train(expert_data, epochs=10, expert_batch_size=32, generator_steps=2048): + env = NRasterized(**env_settings) + + 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(NRasterized, n_envs=1, env_kwargs=env_settings) + discriminator = adversarial.GAIL( + expert_data=expert_data, + expert_batch_size=expert_batch_size, + 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 range(epochs): + train_discriminator(env, generator, discriminator, num_samples=expert_batch_size) + train_generator(env, generator, discriminator, num_samples=generator_steps) + + return generator + +# %% + +with open("data/NormalizedIntersimpleExpertMu.001_NRasterizedAgent51w36h36mppx2.pkl", "rb") as f: + trajectories = pickle.load(f) +transitions = rollout.flatten_trajectories(trajectories) +generator = train(transitions, epochs=2, expert_batch_size=2, generator_steps=2) + +generator.save(model_name) + +# %% +model = stable_baselines3.PPO.load(model_name) + +env = NRasterized(**env_settings) + +s = env.reset() +done = False +env.render() +while not done: + m = available_actions(env) + ch, _, _ = generator.policy.predict({'obs': torch.tensor(s).unsqueeze(0), 'mask': torch.tensor(m).unsqueeze(0)}) + plan = list(generate_plan(env, ch)) + + while not done and plan and feasible(env, plan): + a, plan = plan[0], plan[1:] + s, _, done, _ = env.step(a) + env.render() + +env.close(filestr='render/'+model_name) + +# %% Tests + +def test_ll_transitions_vs_expert_data(): + pass + +def test_hl_transitions(): + pass