From fc2cd936a8f05a4983c594c554a671bfba38f686 Mon Sep 17 00:00:00 2001 From: ebuehrle <43623224+ebuehrle@users.noreply.github.com> Date: Fri, 29 Oct 2021 14:40:11 +0200 Subject: [PATCH 1/2] Copy options to scratch --- scratch/etienne/intersimple/gail/options2.py | 165 +++++++++++++++++++ 1 file changed, 165 insertions(+) create mode 100644 scratch/etienne/intersimple/gail/options2.py diff --git a/scratch/etienne/intersimple/gail/options2.py b/scratch/etienne/intersimple/gail/options2.py new file mode 100644 index 0000000..b160852 --- /dev/null +++ b/scratch/etienne/intersimple/gail/options2.py @@ -0,0 +1,165 @@ +import gym +import torch +from src.util.collisions import feasible +import numpy as np + +class OptionsEnv(gym.Wrapper): + + def __init__(self, env, options=[(0, 5), (5, 5), (10, 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) + if not self.m.any(): + # action 0 is considered safe fallback + self.m[0] = True + + self.ch, self.value, self.log_prob = generator.policy.forward({ + 'obs': torch.tensor(self.s).unsqueeze(0).to(generator.policy.device), + 'mask': 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, safety_plan(self.env, self.plan)) or self.m.sum() == 1): + + 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.predict_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.cpu(), + 'reward': self.r, + '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 safety_plan(env, plan): + return np.concatenate((plan, np.array(5 * [env._env._min_acc])), axis=0) + +def available_actions(env, options): + """Return mask of available actions given current `env` state.""" + plans = [generate_plan(env, i, options) for i, _ in enumerate(options)] + # is emergency braking still possible? + plans = list(map(lambda p: safety_plan(env, p), plans)) + + T = max(len(p) for p in plans) + plans = [np.pad(p, ((0, T-len(p)),), constant_values=np.nan) for p in plans] + plans = np.stack(plans, axis=0) + + valid = feasible(env, plans) + 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 From 66c10f528058793f98151d8861b7e7037ccfcb21 Mon Sep 17 00:00:00 2001 From: ebuehrle <43623224+ebuehrle@users.noreply.github.com> Date: Fri, 29 Oct 2021 16:36:12 +0200 Subject: [PATCH 2/2] Make options env compatible with PPO --- scratch/etienne/intersimple/gail/options2.py | 188 ++++++++---------- .../gail_options_image_random_location.py | 33 +-- 2 files changed, 96 insertions(+), 125 deletions(-) diff --git a/scratch/etienne/intersimple/gail/options2.py b/scratch/etienne/intersimple/gail/options2.py index b160852..d4ae929 100644 --- a/scratch/etienne/intersimple/gail/options2.py +++ b/scratch/etienne/intersimple/gail/options2.py @@ -2,13 +2,22 @@ import gym import torch from src.util.collisions import feasible import numpy as np +from collections import deque +import itertools + +def imitation_discriminator(discriminator): + return lambda obs, action, next_obs, done: discriminator.discrim_net.predict_reward_train( + state=torch.tensor(obs).unsqueeze(0).to(discriminator.discrim_net.device()), + action=torch.tensor([[action]]).to(discriminator.discrim_net.device()), + next_state=torch.tensor(next_obs).unsqueeze(0).to(discriminator.discrim_net.device()), # unused + done=torch.tensor(done).unsqueeze(0).to(discriminator.discrim_net.device()), # unused + ).item() class OptionsEnv(gym.Wrapper): - - def __init__(self, env, options=[(0, 5), (5, 5), (10, 5)], *args, **kwargs): - """option 0 is treated as safe fallback""" + def __init__(self, env, options, discriminator, ll_buffer_capacity, *args, **kwargs): super().__init__(env, *args, **kwargs) + self.options = options num_hl_options = len(self.options) self.action_space = gym.spaces.Discrete(num_hl_options) @@ -17,118 +26,72 @@ class OptionsEnv(gym.Wrapper): 'mask': gym.spaces.Box(low=0, high=1, shape=(num_hl_options,)), }) - def _after_choice(self): - pass + self.discriminator = discriminator + self.ll_buffer_capacity = ll_buffer_capacity + self.ll_buffer = deque(maxlen=ll_buffer_capacity) - 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) - if not self.m.any(): - # action 0 is considered safe fallback - self.m[0] = True - - self.ch, self.value, self.log_prob = generator.policy.forward({ - 'obs': torch.tensor(self.s).unsqueeze(0).to(generator.policy.device), - 'mask': 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, safety_plan(self.env, self.plan)) or self.m.sum() == 1): - - 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.predict_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.cpu(), - 'reward': self.r, - 'episode_start': self.episode_start, - 'value': self.value.detach(), - 'log_prob': self.log_prob.detach(), - 'done': self.done, + @staticmethod + def _hl_observation(obs, mask): + return { + 'obs': obs, + 'mask': mask, } - def sample_hl(self, policy, discriminator): - self.discriminator = discriminator - return self.sample(policy) + def reset(self): + self.done = False + self.obs = self.env.reset() + self.m = available_actions(self.env, self.options) + return self._hl_observation(self.obs, self.m) + + def _ll_step(self, action): + return self.env.step(action) + + def step(self, action): + assert self.m[action] + assert not self.done -class RenderOptions(LLOptions): + plan = list(map(float, generate_plan(self.env, action, self.options))) + reward = 0 + steps = 0 - def _after_step(self): - super()._after_step() + while not self.done and plan and \ + (feasible(self.env, safety_plan(self.env, plan)) or self.m.sum() == 1): + + a, plan = plan[0], plan[1:] + a = self.env._normalize(a) + + next_obs, _, self.done, info = self._ll_step(a) + + reward += self.discount**steps * self.discriminator(self.obs, a, next_obs, self.done) + + self.ll_buffer.append({ + 'obs': self.obs, + 'next_obs': next_obs, + 'acts': np.array((a,)), + 'dones': np.array(self.done), + }) + + steps += 1 + self.obs = next_obs + + self.m = available_actions(self.env, self.options) + + return self._hl_observation(self.obs, self.m), reward, self.done, info + + def sample_ll(self, n): + assert n <= self.ll_buffer_capacity, f'Sample size of {n} exceeds buffer capacity of {self.ll_buffer_capacity}' + assert n <= len(self.ll_buffer), f'Sample size of {n} exceeds buffer size of {len(self.ll_buffer)}' + return list(itertools.islice(self.ll_buffer, n)) + +class RenderOptions(OptionsEnv): + + def __init__(self, options, *args, **kwargs): + super().__init__(options, discriminator=lambda s, a, n, d: 0, ll_buffer_capacity=0, *args, **kwargs) + + def _ll_step(self): + out = super()._ll_step() self.env.render() + return out def close(self, *args, **kwargs): self.env.close(*args, **kwargs) @@ -137,7 +100,9 @@ def safety_plan(env, plan): return np.concatenate((plan, np.array(5 * [env._env._min_acc])), axis=0) def available_actions(env, options): - """Return mask of available actions given current `env` state.""" + """Return mask of available actions given current `env` state. + Action 0 is considered safe fallback. + """ plans = [generate_plan(env, i, options) for i, _ in enumerate(options)] # is emergency braking still possible? plans = list(map(lambda p: safety_plan(env, p), plans)) @@ -147,6 +112,9 @@ def available_actions(env, options): plans = np.stack(plans, axis=0) valid = feasible(env, plans) + if not valid.any(): + valid[0] = True + return valid def target_velocity_plan(current_v: float, target_v: float, t: int, dt: float): diff --git a/scratch/etienne/intersimple/gail_options_image_random_location.py b/scratch/etienne/intersimple/gail_options_image_random_location.py index 3a3f34c..330dedb 100644 --- a/scratch/etienne/intersimple/gail_options_image_random_location.py +++ b/scratch/etienne/intersimple/gail_options_image_random_location.py @@ -20,8 +20,9 @@ from imitation.util import logger from stable_baselines3.common.env_util import make_vec_env from tqdm import tqdm from src.policies.options import OptionsCnnPolicy -from src.gail.options import OptionsEnv, LLOptions, HLOptions, RenderOptions -from src.gail.train import train_discriminator, train_generator +from src.gail.train import flatten_transitions + +from gail.options2 import OptionsEnv, RenderOptions, imitation_discriminator model_name = 'gail_options_image_random_location' env_settings = {'width': 70, 'height': 70, 'm_per_px': 1, 'map_color': 128, 'mu': 0.001} @@ -30,12 +31,13 @@ ALL_OPTIONS = [(v,t) for v in [0,2,4,6,8] for t in [5, 10, 20]] # option 0 is sa def train( expert_data, - epochs=200, expert_batch_size=256, - generator_steps=1024, - discount=0.99, n_disc_updates_per_round=2, + generator_steps=512, + generator_total_steps=2048, n_gen_updates_per_round=10, + discount=0.99, + epochs=100, ): env = NRasterizedRouteRandomAgentLocation(**env_settings) env.discount = discount @@ -55,24 +57,25 @@ def train( gen_algo=stable_baselines3.PPO("CnnPolicy", venv), # unused ) + options_env = OptionsEnv(env, discriminator=imitation_discriminator(discriminator), options=ALL_OPTIONS, ll_buffer_capacity=expert_batch_size) generator = stable_baselines3.PPO( OptionsCnnPolicy, - OptionsEnv(env, options=ALL_OPTIONS), + options_env, verbose=1, n_steps=generator_steps, n_epochs=n_gen_updates_per_round, ) - # 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, options=ALL_OPTIONS), generator, discriminator, num_samples=expert_batch_size, n_updates=n_disc_updates_per_round) - train_generator(HLOptions(env, options=ALL_OPTIONS), generator, discriminator, num_samples=generator_steps) + # train generator + generator.learn(total_timesteps=generator_total_steps) + + # train discriminator + generator_samples = options_env.sample_ll(expert_batch_size) + generator_samples = flatten_transitions(generator_samples) + for _ in range(n_disc_updates_per_round): + discriminator.train_disc(gen_samples=generator_samples) + generator.save(model_name) return generator