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] 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