diff --git a/scratch/etienne/intersimple/gail/options2.py b/scratch/etienne/intersimple/gail/options2.py new file mode 100644 index 0000000..d4ae929 --- /dev/null +++ b/scratch/etienne/intersimple/gail/options2.py @@ -0,0 +1,133 @@ +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, 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) + self.observation_space = gym.spaces.Dict({ + 'obs': env.observation_space, + 'mask': gym.spaces.Box(low=0, high=1, shape=(num_hl_options,)), + }) + + self.discriminator = discriminator + self.ll_buffer_capacity = ll_buffer_capacity + self.ll_buffer = deque(maxlen=ll_buffer_capacity) + + @staticmethod + def _hl_observation(obs, mask): + return { + 'obs': obs, + 'mask': mask, + } + + 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 + + plan = list(map(float, generate_plan(self.env, action, self.options))) + reward = 0 + steps = 0 + + 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) + +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. + 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)) + + 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) + if not valid.any(): + valid[0] = True + + 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_options_image_random_location.py b/scratch/etienne/intersimple/gail_options_image_random_location.py index 750ca7d..fc4909b 100644 --- a/scratch/etienne/intersimple/gail_options_image_random_location.py +++ b/scratch/etienne/intersimple/gail_options_image_random_location.py @@ -5,13 +5,7 @@ sys.path.append('../../../') from src.discriminator import CnnDiscriminatorFlatAction from imitation.algorithms import adversarial import stable_baselines3 -import torch.utils.data -import numpy as np from intersim.envs import NRasterizedRouteSpeedRandomAgentLocation -import itertools -from torch.distributions import Categorical -import gym -import torch import pickle import imitation.data.rollout as rollout import tempfile @@ -20,8 +14,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 +25,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=1024, - generator_steps=1024, + discriminator_updates_per_round=10, + generator_steps=256, + generator_total_steps=1024, + generator_updates_per_round=10, discount=0.99, - n_disc_updates_per_round=10, - n_gen_updates_per_round=10, + epochs=200, ): env = NRasterizedRouteSpeedRandomAgentLocation(**env_settings) env.discount = discount @@ -55,24 +51,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, + n_epochs=generator_updates_per_round, ) 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(discriminator_updates_per_round): + discriminator.train_disc(gen_samples=generator_samples) + generator.save(model_name) return generator