From dcf8212028930c9c5876b129d8393e1b564066e2 Mon Sep 17 00:00:00 2001 From: Arec Date: Wed, 20 Oct 2021 12:28:12 -0700 Subject: [PATCH 1/7] adding comments to gail_options_image, combining environments for options gail, and fixing bug where last state is yielded in hl buffer --- .../arec/intersimple/gail_options_image.py | 143 ++++- scratch/arec/intersimple/options_gail.py | 510 ++++++++++++++++++ 2 files changed, 644 insertions(+), 9 deletions(-) create mode 100644 scratch/arec/intersimple/options_gail.py diff --git a/scratch/arec/intersimple/gail_options_image.py b/scratch/arec/intersimple/gail_options_image.py index 26ef3af..b347c2d 100644 --- a/scratch/arec/intersimple/gail_options_image.py +++ b/scratch/arec/intersimple/gail_options_image.py @@ -23,17 +23,38 @@ logging.basicConfig(level=logging.DEBUG) 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): - + """ + Class for high-level options policy (generator) + """ 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) + """ + Return prior distribution over high-level options (before masking) + Args: + s (torch.tensor): observation + Returns: + values (torch.tensor): values from critic + dist (torch.distributions): prior distribution over actions + """ + 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): + """ + Will mask invalid states before making action selections + Args: + obs: dict with keys: + obs (torch.tensor): (B,o) true observations + mask (torch.tensor): (B,m) mask over valid actions + Returns: + ch (torch.tensor): (B,a) sampled actions + values (torch.tensor): (B,) predicted value at observation + log_probs (torch.tensor): (B,) log probabilities of selected actions + """ s, m = obs['obs'], obs['mask'] values, prior = self._prior_distribution(s) posterior = Categorical(prior.probs * m) @@ -41,14 +62,31 @@ class OptionsCnnPolicy(stable_baselines3.common.policies.ActorCriticCnnPolicy): return ch, values, posterior.log_prob(ch) def evaluate_actions(self, obs, ch): + """ + Evaluate particular actions + Args: + obs: dict with keys: + obs (torch.tensor): (B,o) true observations + mask (torch.tensor): (B,m) masks over valid actions + ch (torch.tensor): (B,a) selected actions + Returns: + values (torch.tensor): (B,) predicted value at observation + log_probs (torch.tensor): (B,) log probabilities of selected actions + ent (torch.tensor): (B,) entropy of each distribution over actions + """ 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 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) @@ -67,51 +105,88 @@ 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 + if self.done: + # reset environment self.s = self.env.reset() self.m = available_actions(self.env) self.done = False self.episode_start = True + # set the action, the value of the start state, and the logprob of the action + # according to the current environment state and mask 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), }) + + # store a float list of actions to take given the option selected in the environment self.plan = list(map(float, generate_plan(self.env, self.ch))) + # run whatever _after_choice might dictate in a child class self._after_choice() + # some checks assert not self.done assert self.plan assert feasible(self.env, self.plan, self.ch) + # execute the option so long as the episode isn't complete and the plan is still feasible while not self.done and self.plan and feasible(self.env, self.plan, self.ch): + + # pop first action self.a, self.plan = self.plan[0], self.plan[1:] + + # normalize action ?? self.a = self.env._normalize(self.a) + + # step through environment self.nexts, _, self.done, _ = self.env.step(self.a) self.nextm = available_actions(self.env) + # run whatever _after_step might dictate in child class self._after_step() + # update state and mask to current self.s = self.nexts self.m = self.nextm - + + # transitions yielded from self._transitions() functions specied in child classes yield from self._transitions() + ### NOTE: only yields after a full option has been executed / exited + 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) - self.observation_space = self.observation_space['obs'] + # 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, @@ -120,9 +195,17 @@ class LLOptions(OptionsEnv): }) def _transitions(self): + """ + Yield from the transition buffer + """ yield from self._transition_buffer def sample_ll(self, policy): + """ + Not quite sure how this works???? + Why would you do this over LLOptions.sample(policy) + """ + # What happens if you return a yield from ???????? return self.sample(policy) class HLOptions(OptionsEnv): @@ -132,10 +215,16 @@ class HLOptions(OptionsEnv): super().__init__(*args, **kwargs) def _after_choice(self): + """ + After an option selection, initialize total reward and number of steps + """ 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()), @@ -145,6 +234,18 @@ 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': {'obs': self.s, 'mask': self.m}, 'action': self.ch, @@ -156,16 +257,29 @@ class HLOptions(OptionsEnv): } def sample_hl(self, policy, discriminator): + """ + Args: + policy + discriminator: function with which to score rewards + Returns: + gen: an 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): @@ -307,6 +421,18 @@ 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) env.discount = discount @@ -367,13 +493,12 @@ if __name__ == '__main__': discount=0.99 ) - generator.save(model_name) + generator.save(model_name) # save ppo sb3 generator class # %% - model = stable_baselines3.PPO.load(model_name) + model = stable_baselines3.PPO.load(model_name) # not actually used env = RenderOptions(NRasterizedRandomAgent(**env_settings)) - for s in env.sample_ll(generator): if s['dones']: break diff --git a/scratch/arec/intersimple/options_gail.py b/scratch/arec/intersimple/options_gail.py new file mode 100644 index 0000000..4ab8f12 --- /dev/null +++ b/scratch/arec/intersimple/options_gail.py @@ -0,0 +1,510 @@ +# %% +from gail.discriminator import CnnDiscriminator, CnnDiscriminatorFlatAction +from imitation.algorithms import adversarial +import stable_baselines3 +import torch.utils.data +import numpy as np +from intersim.envs.intersimple import NRasterized, NRasterizedRandomAgent +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 + +import logging +logging.basicConfig(level=logging.DEBUG) + +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): + """ + Class for high-level options policy (generator) + """ + def __init__(self, observation_space, *args, **kwargs): + super().__init__(observation_space['obs'], *args, **kwargs) + + def _prior_distribution(self, s): + """ + Return prior distribution over high-level options (before masking) + Args: + s (torch.tensor): observation + Returns: + values (torch.tensor): values from critic + dist (torch.distributions): prior distribution over actions + """ + 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): + """ + Will mask invalid states before making action selections + Args: + obs: dict with keys: + obs (torch.tensor): (B,o) true observations + mask (torch.tensor): (B,m) mask over valid actions + Returns: + ch (torch.tensor): (B,a) sampled actions + values (torch.tensor): (B,) predicted value at observation + log_probs (torch.tensor): (B,) log probabilities of selected actions + """ + 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): + """ + Evaluate particular actions + Args: + obs: dict with keys: + obs (torch.tensor): (B,o) true observations + mask (torch.tensor): (B,m) masks over valid actions + ch (torch.tensor): (B,a) selected actions + Returns: + values (torch.tensor): (B,) predicted value at observation + log_probs (torch.tensor): (B,) log probabilities of selected actions + ent (torch.tensor): (B,) entropy of each distribution over actions + """ + 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 + +class OptionsEnv(gym.Wrapper): + """ + Wrap an intersimple environment with an options generator + """ + def __init__(self, env, render=False, *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,)), + }) + self._hl_transition_buffer = [] + self._ll_transition_buffer = [] + self.render=render + + def _after_option_choice(self): + """ + After initial option choice, + """ + self._hl_r = 0 + self._hl_steps = 0 + + def _after_step(self): + """ + After each step, add the ll transition to the appropriate buffer, add to reward, add to steps, and possibly render + """ + + self._ll_transition_buffer.append({ + 'obs': self.s, + 'next_obs': self.nexts, + 'acts': np.array((self.a,)), + 'dones': np.array(self.done), + }) + 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 + if self.render: + self.env.render() + + def _after_option(self): + """ + After each low-level action, add the discounted discriminated reward score (given a discriminator) + """ + self._hl_transition_buffer.append({ + 'obs': {'obs': self.os, 'mask': self.m}, + '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 close(self, *args, **kwargs): + """ + On 'close', close the environment + """ + self.env.close(*args, **kwargs) + + def sample(self, generator, controller): + """ + yield transitions using a generator + Args: + generator (sb3.PPO) + controller (str): 'high' or 'low' to yield from proper buffer + Yields: + + """ + self.done = True + # DO I WANT TO EMPTY THE BUFFERS??? Probs naw + while True: + + # yield from buffers to empty what was stored previously + if controller = 'high': + yield from self._hl_transition_buffer + elif controller == 'low': + yield from self._ll_transition_buffer + else: + raise('Improper buffer') + + self.episode_start = False + if self.done: + # reset environment + self.s = self.env.reset() + self.done = False + self.episode_start = True + + self.os = self.s.copy() # option start state + self.m = available_actions(self.env) + + # set the action, the value of the start state, and the logprob of the action + # according to the current environment state and mask + self.ch, self.value, self.log_prob = generator.policy.predict({ + 'obs': torch.tensor(self.os).unsqueeze(0).to(generator.policy.device), + 'mask': torch.tensor(self.m).unsqueeze(0).to(generator.policy.device), + }) + + # store a float list of actions to take given the option selected in the environment + self.plan = list(map(float, generate_plan(self.env, self.ch))) + + # run whatever _after_choice might dictate in a child class + self._after_option_choice() + + # some checks + assert not self.done + assert self.plan + assert feasible(self.env, self.plan, self.ch) + + # execute the option so long as the episode isn't complete and the plan is still feasible + while not self.done and self.plan and feasible(self.env, self.plan, self.ch): + + # pop first action + self.a, self.plan = self.plan[0], self.plan[1:] + + # normalize action ?? + self.a = self.env._normalize(self.a) + + # step through environment + self.nexts, _, self.done, _ = self.env.step(self.a) + + # run whatever _after_step might dictate in child class + self._after_step() + + # update state and mask to current + self.s = self.nexts + + # run whatever to do after option + self._after_option() + + def sample_ll(self, policy): + """ + Not quite sure how this works???? + Why would you do this over LLOptions.sample(policy) + """ + return self.sample(policy, 'low') + + def sample_hl(self, policy, discriminator): + """ + Args: + policy + discriminator: function with which to score rewards + Returns: + gen: an which samples high-level transitions from the environment + """ + self.discriminator = discriminator + return self.sample(policy) + +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`.""" + 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 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 check_future_collisions_circles(env, actions, n_circles:int=2): + """Checks whether `env._agent` would collide with other agents assuming `actions` as input. + + Vehicles are (over-)approximated by multiple 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 + """ + assert n_circles >= 2 + 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) + centers = states[:, :, :, :2] + psi = states[:, :, :, 3] + lon = torch.stack([psi.cos(), psi.sin()],dim=-1) # (B, T, nv, 2) + + # offset between [-env._env.lengths+env._env.widths/2, env._env.lengths/2-env._env.widths/2] + back = (-env._env._lengths/2+env._env._widths/2).unsqueeze(-1) # (nv, 1) + length = (env._env._lengths-env._env._widths).unsqueeze(-1) # (nv, 1) + diff_d = back + length*(torch.arange(n_circles)/(n_circles-1)).unsqueeze(0) # (nv, n_circles) + assert diff_d.shape == (nv, n_circles) + + offsets = diff_d[None, None, :, :, None] * lon[:, :, :, None, :] + assert offsets.shape == (B, T, nv, n_circles, 2) + + expanded_centers=centers.unsqueeze(-2) + offsets #(B, T, nv, n_circles, 2) + assert expanded_centers.shape == (B, T, nv, n_circles, 2) + agent_centers = expanded_centers[:,:,env._agent:env._agent+1,:,:] #(B, T, 1, n_circles, 2) + ds = expanded_centers.reshape((B, T, nv*n_circles, 1, 2)) - agent_centers #(B, T, nv*nc,1, 2) - (B, T, 1, nc, 2) = (B, T, nv*nc, nc, 2) + + distance = (ds**2).sum(-1).sqrt().reshape((B, T, nv, n_circles, n_circles)) # (B, T, nv, nc, nc) + 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, n_circles, n_circles) + + radius = env._env._widths*np.sqrt(2) / 2 + min_distance = radius[env._agent] + radius + min_distance = min_distance[None, None, :, None, None] + assert min_distance.shape == (1, 1, nv, 1, 1) + + return (distance > min_distance).all(-1).all(-1).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.""" + + # 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 + valid = check_future_collisions_circles(env, [full_plan]) + return ch == 0 or valid.item() + +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} + + with open("data/NormalizedIntersimpleExpertMu.001_NRasterizedIncrementingAgentw36h36mppx2.pkl", "rb") as f: + trajectories = pickle.load(f) + #import pdb + #pdb.set_trace() + transitions = rollout.flatten_trajectories(trajectories) + generator = train( + transitions, + env_class=env_class, + env_settings=env_settings, + epochs=2, + discrim_batch_size=32, + generator_steps=2048, + discount=0.99 + ) + + generator.save(model_name) # save ppo sb3 generator class + + # %% + model = stable_baselines3.PPO.load(model_name) # not actually used + + env = OptionsGail(NRasterizedRandomAgent(**env_settings), render=True) + for s in env.sample_ll(generator): + 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 From da1fb112692748a2a264965a1fcb110a0298cfb1 Mon Sep 17 00:00:00 2001 From: Arec Date: Thu, 21 Oct 2021 05:21:06 -0700 Subject: [PATCH 2/7] adding functions to process expert data across locations and tracks in intersimple environment --- generate_demos.sh | 11 ++ src/__init__.py | 4 +- src/advil/__init__.py | 0 src/data/__init__.py | 1 + src/{ => data}/data_utils.py | 0 src/data/expert.py | 214 ++++++++++++++++++++++++++++++++++ src/{ => data}/expert_data.py | 0 7 files changed, 228 insertions(+), 2 deletions(-) create mode 100755 generate_demos.sh delete mode 100644 src/advil/__init__.py create mode 100644 src/data/__init__.py rename src/{ => data}/data_utils.py (100%) create mode 100644 src/data/expert.py rename src/{ => data}/expert_data.py (100%) diff --git a/generate_demos.sh b/generate_demos.sh new file mode 100755 index 0000000..447068a --- /dev/null +++ b/generate_demos.sh @@ -0,0 +1,11 @@ +#DEFAULT PARAMETERS: +# locs:list=None, (default to all locations) +# tracks:list=None, (default to all tracks) +# env_class:str='NRasterizedIncrementingAgent', +# env_args:dict={width:36,height:36,m_per_px:2}, +# expert_class:str='NormalizedIntersimpleExpert', +# expert_args:dict={mu:0.001}): + +cd ./src/data +python -m expert --locs='[DR_USA_Roundabout_FT]' +cd ../ \ No newline at end of file diff --git a/src/__init__.py b/src/__init__.py index daabc89..d44975d 100644 --- a/src/__init__.py +++ b/src/__init__.py @@ -1,3 +1,3 @@ -from src.expert_data import generate_expert_data, load_expert_data -from src.data_utils import InteractionDatasetSingleAgent +from src.data.expert_data import generate_expert_data, load_expert_data +from src.data.data_utils import InteractionDatasetSingleAgent from src.metrics import metrics \ No newline at end of file diff --git a/src/advil/__init__.py b/src/advil/__init__.py deleted file mode 100644 index e69de29..0000000 diff --git a/src/data/__init__.py b/src/data/__init__.py new file mode 100644 index 0000000..49d78bc --- /dev/null +++ b/src/data/__init__.py @@ -0,0 +1 @@ +from src.data.expert_trajectories import demonstrations, load_experts, process_experts \ No newline at end of file diff --git a/src/data_utils.py b/src/data/data_utils.py similarity index 100% rename from src/data_utils.py rename to src/data/data_utils.py diff --git a/src/data/expert.py b/src/data/expert.py new file mode 100644 index 0000000..fb6b754 --- /dev/null +++ b/src/data/expert.py @@ -0,0 +1,214 @@ +from intersim.envs.intersimple import Intersimple +from stable_baselines3.common.policies import BasePolicy +import gym +import intersim.envs.intersimple +import pickle +from tqdm import tqdm +import imitation.data.rollout as rollout +from stable_baselines3.common.vec_env.dummy_vec_env import DummyVecEnv +from imitation.data.wrappers import RolloutInfoWrapper +import copy +import os + +class IntersimExpert(BasePolicy): + + def __init__(self, intersim_env, mu=0, *args, **kwargs): + super().__init__( + observation_space=gym.spaces.Space(), + action_space=gym.spaces.Space(), + *args, **kwargs + ) + self._intersim = intersim_env + self._mu = mu + + def forward(self, *args, **kwargs): + raise NotImplementedError() + + def _predict(self, *args, **kwargs): + raise NotImplementedError() + + def _action(self): + target_t = min(self._intersim._ind + 1, len(self._intersim._svt.simstate) - 1) + target_state = self._intersim._svt.simstate[target_t] + return self._intersim.target_state(target_state, mu=self._mu) + + def predict(self, *args, **kwargs): + return self._action(), None + +class IntersimpleExpert(BasePolicy): + + def __init__(self, intersimple_env, mu=0, *args, **kwargs): + super().__init__( + observation_space=intersimple_env.observation_space, + action_space=intersimple_env.action_space, + *args, **kwargs + ) + self._intersimple = intersimple_env + self._intersim_expert = IntersimExpert(intersimple_env._env, mu=mu) + + def forward(self, *args, **kwargs): + raise NotImplementedError() + + def _predict(self, *args, **kwargs): + raise NotImplementedError() + + def _action(self): + return self._intersim_expert._action()[self._intersimple._agent] + + def predict(self, *args, **kwargs): + return self._action(), None + +class NormalizedIntersimpleExpert(IntersimpleExpert): + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + + def predict(self, *args, **kwargs): + action, _ = super().predict(*args, **kwargs) + return self._intersimple._normalize(action), None + +class DummyVecEnvPolicy(BasePolicy): + + def __init__(self, experts): + self._experts = [e() for e in experts] + + def forward(self, *args, **kwargs): + raise NotImplementedError() + + def _predict(self, *args, **kwargs): + raise NotImplementedError() + + def predict(self, *args, **kwargs): + predictions = [e.predict() for e in self._experts] + actions = [p[0] for p in predictions] + states = [p[1] for p in predictions] + return actions, states + + def forward(self, *args, **kwargs): + raise NotImplementedError() + + def _predict(self, *args, **kwargs): + raise NotImplementedError() + +def save_video(env, expert): + env.reset() + env.render() + done = False + while not done: + actions, _ = expert.predict() + _, _, done, _ = env.step(actions) + env.render() + env.close() + +def load_experts(expert_files=[]): + """ + Load expert trajectories from files and combine their transitions into a single RB + + Args: + expert_files (list): list of expert file strings + Returns: + transitions (list): list of combined expert episode transitions + """ + transitions = [] + for file in tqdm(expert_files): + with open(file, "rb") as f: + trajectories = pickle.load(f) + transitions = transitions + rollout.flatten_trajectories(trajectories) + return transitions + +def demonstrations(expert='NormalizedIntersimpleExpert', env='NRasterizedIncrementingAgent', path=None, min_timesteps=None, min_episodes=None, video=False, env_args={}, policy_args={}): + """Rollout and save expert demos. + + Usage: + python -m intersimple.expert + Args: + expert (class): class of expert + env (class): class of env intersim.envs.intersimple + path (str): path to store output + min_timesteps (int): min number of timesteps for call to rollout.rollout_and_save + min_episodes (int): min number of episodes for call to rollout.rollout_and_save + video (bool): whether to save a video of the expert until a single environment instantiation stops + env_args (dict): dictionary of kwargs when instantiating environment class + policy_args (dict): dictionary of kwargs when instantiating Expert policy + """ + + Env = intersim.envs.intersimple.__dict__[env] + Expert = globals()[expert] + + env = Env(**env_args) + info_env = RolloutInfoWrapper(env) # getting rollout info (dictionary) from environment + venv = DummyVecEnv([lambda: info_env]) # making a DummyVecEnv with a list of a function that when called returns the rollout info + + policy = Expert(env, **policy_args) # instantiate an expert policy from specified class with instantiated environment and policy kwargs + venv_policy = DummyVecEnvPolicy([lambda: policy]) # make a DummyVecEnvPolicy with a list of a function that when called returns the Expert policy + + if min_timesteps is None and min_episodes is None: + min_episodes = env.nv # one episode per vehicle being controlled in environment (hopefully an incrementing agent environment) + + if video: + save_video(env, policy) + + path = path or (policy.__class__.__name__ + '_' + env.__class__.__name__ + '.pkl') + suntil = rollout.make_sample_until( + min_timesteps=min_timesteps, + min_episodes=min_episodes, + ) + rollout.rollout_and_save( + path=path, + policy=venv_policy, + venv=venv, + sample_until=suntil + ) + +def process_experts(filename:str='expert.pkl', + locs:list=None, + tracks:list=None, + env_class:str='NRasterizedIncrementingAgent', + env_args:dict={'width':36,'height':36,'m_per_px':2}, + expert_class:str='NormalizedIntersimpleExpert', + expert_args:dict={'mu':0.001}): + """ + Process all experts in the Interaction Dataset + For now, using NormalizedIntersimpleExpert with NRasterizedIncrementingAgent environment + + Args: + filename (str): name for track file + locs (list): list of location ids + tracks (list): list of track numbers + env_class (str): class of environment + env_args (dict): default environment kwargs + expert_class (str): class of expert + expert_args (dict): default expert kwargs + """ + locs = locs or intersim.LOCATIONS + tracks = tracks or range(intersim.MAX_TRACKS) + pbar = tqdm(total=len(locs)*len(tracks)) + for loc in locs: + for track in tracks: + + iloc = intersim.LOCATIONS.index(loc) + + it_env_args = copy.deepcopy(env_args) + it_env_args.update({ + 'loc':iloc, + 'track':track, + }) + out_folder = os.path.join('expert_data',loc, 'track%04i'%(track)) + if not os.path.isdir(out_folder): + os.makedirs(out_folder) + it_path = os.path.join(out_folder,filename) + + demonstrations( + expert=expert_class, + env=env_class, + path=it_path, + env_args=it_env_args, + policy_args=expert_args, + ) + pbar.update(1) + pbar.close() + +if __name__=='__main__': + import fire + fire.Fire(process_experts) + diff --git a/src/expert_data.py b/src/data/expert_data.py similarity index 100% rename from src/expert_data.py rename to src/data/expert_data.py From 45a99978e421d3828bf96c91286e480e082b70a5 Mon Sep 17 00:00:00 2001 From: Arec Date: Thu, 21 Oct 2021 05:33:05 -0700 Subject: [PATCH 3/7] adding discriminators to main folder, utilities to render a video from a saved model --- src/discriminator/__init__.py | 1 + src/discriminator/discriminator.py | 101 +++++++++++++++++++++++++++++ src/util/__init__.py | 1 + src/util/render_env.py | 58 +++++++++++++++++ tests/test_discriminator.py | 45 +++++++++++++ 5 files changed, 206 insertions(+) create mode 100644 src/discriminator/__init__.py create mode 100644 src/discriminator/discriminator.py create mode 100644 src/util/render_env.py create mode 100644 tests/test_discriminator.py diff --git a/src/discriminator/__init__.py b/src/discriminator/__init__.py new file mode 100644 index 0000000..4ad66b2 --- /dev/null +++ b/src/discriminator/__init__.py @@ -0,0 +1 @@ +from src.discriminator.discriminator import * \ No newline at end of file diff --git a/src/discriminator/discriminator.py b/src/discriminator/discriminator.py new file mode 100644 index 0000000..1c9664c --- /dev/null +++ b/src/discriminator/discriminator.py @@ -0,0 +1,101 @@ +import torch + +# imitation.rewards.discrim_nets.DiscrimNetGAIL is composed of self.discriminator (nn.Module), +# which gets called with inputs (state, action) when needed. + +class CnnDiscriminator(torch.nn.Module): + """ConvNet similar to stable_baselines3.common.policies.ActorCriticCnnPolicy.""" + + def __init__(self, env): + super().__init__() + + obs_channels, _, _ = env.observation_space.shape + (action_size,) = env.action_space.shape + in_channels = obs_channels + action_size + + self.cnn = torch.nn.Sequential( + torch.nn.Conv2d(in_channels, 32, kernel_size=(8, 8), stride=(4, 4)), # 5+1 -> 32 + torch.nn.ReLU(), + torch.nn.Conv2d(32, 64, kernel_size=(4, 4), stride=(2, 2)), # 32 -> 64 + torch.nn.ReLU(), + torch.nn.Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1)), # 64 -> 64 + torch.nn.ReLU(), + torch.nn.Flatten(start_dim=1, end_dim=-1), + torch.nn.LazyLinear(512), # 28224 -> 512 + torch.nn.ReLU(), + torch.nn.LazyLinear(1), # 512 -> 1 + ) + + @staticmethod + def _concatenate(state, action): + b, _, h, w = state.shape + _, a = action.shape + act = action.unsqueeze(-1).unsqueeze(-1).expand((b, a, h, w)) + sa = torch.cat((state, act), -3) + return sa + + def forward(self, state, action): + sa = self._concatenate(state, action) + assert sa.ndim == 4 + return self.cnn(sa).squeeze(1) + +class CnnDiscriminatorFlatAction(torch.nn.Module): + """ConvNet similar to stable_baselines3.common.policies.ActorCriticCnnPolicy.""" + + def __init__(self, env): + super().__init__() + + obs_channels, _, _ = env.observation_space.shape + (action_size,) = env.action_space.shape + in_channels = obs_channels + + self.cnn = torch.nn.Sequential( + torch.nn.Conv2d(in_channels, 32, kernel_size=(8, 8), stride=(4, 4)), # in_channels -> 32 + torch.nn.ReLU(), + torch.nn.Conv2d(32, 64, kernel_size=(4, 4), stride=(2, 2)), # 32 -> 64 + torch.nn.ReLU(), + torch.nn.Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1)), # 64 -> 64 + torch.nn.ReLU(), + torch.nn.Flatten(start_dim=1, end_dim=-1), + torch.nn.LazyLinear(128), # 28224 -> 128 + ) + self.decoder = torch.nn.Sequential( + torch.nn.LazyLinear(64), #128 + 2 -> 64 + torch.nn.ReLU(), + torch.nn.LazyLinear(64), #64 -> 64 + torch.nn.ReLU(), + torch.nn.LazyLinear(1) #64 -> 1 + ) + + @staticmethod + def _concatenate(state, action): + b, s= state.shape + b, a = action.shape + sa = torch.cat((state, action), -1) + return sa + + def forward(self, state, action): + s = self.cnn(state.float()) + sa = self._concatenate(s, action) + assert sa.ndim == 2 + return self.decoder(sa).squeeze(1) + +class MlpDiscriminator(torch.nn.Module): + """MLP similar to stable_baselines3.common.policies.ActorCriticPolicy.""" + + def __init__(self, env=None): + super().__init__() + self.flatten = torch.nn.Flatten(start_dim=1, end_dim=-1) + self.mlp = torch.nn.Sequential( + torch.nn.LazyLinear(64), # 42 -> 64 + torch.nn.Tanh(), + torch.nn.LazyLinear(64), # 64 -> 64 + torch.nn.Tanh(), + torch.nn.LazyLinear(1), # 64 -> 1 + ) + + def forward(self, state, action): + flat = self.flatten(state) + sa = torch.cat((action, flat), -1) + assert sa.ndim == 2 + return self.mlp(sa).squeeze(1) diff --git a/src/util/__init__.py b/src/util/__init__.py index e69de29..000365b 100644 --- a/src/util/__init__.py +++ b/src/util/__init__.py @@ -0,0 +1 @@ +from src.util.render_env import * \ No newline at end of file diff --git a/src/util/render_env.py b/src/util/render_env.py new file mode 100644 index 0000000..8223c73 --- /dev/null +++ b/src/util/render_env.py @@ -0,0 +1,58 @@ + +import stable_baselines3 as sb3 +from intersim.envs.intersimple import NRasterized + + +def render_env(model_name='gail_image_multiagent_nocollision', agent=51, environment=NRasterized): + """ + Render a video from an model, agent, and environment + Args: + model_name (str): name of the model + agent (int): agent to start the video from + environment (gym.Env): gym environment class to render environment on + """ + + model = sb3.PPO.load(model_name) + + env = environment(stop_on_collision=False, width=36, height=36, m_per_px=2, agent=agent) + + obs = env.reset() + i=0 + while True and i < 600: + i+=1 + action, _states = model.predict(obs) + obs, rewards, done, info = env.step(action) + env.render(mode='post') + if done: + break + + env.close(filestr='render/'+model_name+'_agent%i'%(agent)) + +def render_options_env(model_name='gail_image_multiagent_nocollision', agent=51, environment=NRasterized): + """ + Render a video from an model, agent, and environment + Args: + model_name (str): name of the model + agent (int): agent to start the video from + environment (gym.Env): gym environment class to render environment on + """ + + model = sb3.PPO.load(model_name) + + env = environment(stop_on_collision=False, width=36, height=36, m_per_px=2, agent=agent) + + obs = env.reset() + i=0 + while True and i < 600: + i+=1 + action, _states = model.predict(obs) + obs, rewards, done, info = env.step(action) + env.render(mode='post') + if done: + break + + env.close(filestr='render/'+model_name+'_agent%i'%(agent)) + +if __name__ == '__main__': + import fire + fire.Fire(render_env) \ No newline at end of file diff --git a/tests/test_discriminator.py b/tests/test_discriminator.py new file mode 100644 index 0000000..1172aa3 --- /dev/null +++ b/tests/test_discriminator.py @@ -0,0 +1,45 @@ +from intersim.envs.intersimple import NRasterized +from src.discriminator import CnnDiscriminator +import torch + +def test_image_concatenation(): + env = NRasterized() + disc = CnnDiscriminator(env) + s = torch.tensor(env.reset()).unsqueeze(0) + a = torch.tensor([[0.5]]) + sa = disc._concatenate(s, a) + + assert s.shape == (1, 5, 200, 200) + assert a.shape == (1, 1) + assert sa.shape == (1, 6, 200, 200) + assert torch.allclose(sa[:, :5], 1.0 * s) + assert (sa[:, 5] == a.unsqueeze(-1)).all() + +def test_image_concatenation3(): + env = NRasterized() + disc = CnnDiscriminator(env) + + s1 = env.reset() + a1 = 0.15 + s2, _, _, _ = env.step(0.9) + a2 = 0.25 + s3, _, _, _ = env.step(-0.9) + a3 = 0.35 + + s = torch.stack([ + torch.tensor(s1), + torch.tensor(s2), + torch.tensor(s3) + ], axis=0) + a = torch.tensor([ + [a1], + [a2], + [a3], + ]) + sa = disc._concatenate(s, a) + + assert s.shape == (3, 5, 200, 200) + assert a.shape == (3, 1) + assert sa.shape == (3, 6, 200, 200) + assert torch.allclose(sa[:, :5], 1.0 * s) + assert (sa[:, 5] == a.unsqueeze(-1)).all() From 05b31092f4d92ace27c8c3f47a5d74802d40be6f Mon Sep 17 00:00:00 2001 From: Arec Date: Thu, 21 Oct 2021 05:59:28 -0700 Subject: [PATCH 4/7] moving options policy to policies, commenting options image, and making the calls to train more flexible --- .../etienne/intersimple/gail_options_image.py | 162 ++++++++++++------ src/policies/__init__.py | 1 + src/policies/options.py | 59 +++++++ 3 files changed, 174 insertions(+), 48 deletions(-) create mode 100644 src/policies/options.py diff --git a/scratch/etienne/intersimple/gail_options_image.py b/scratch/etienne/intersimple/gail_options_image.py index fb2329f..150afe3 100644 --- a/scratch/etienne/intersimple/gail_options_image.py +++ b/scratch/etienne/intersimple/gail_options_image.py @@ -1,55 +1,38 @@ # %% -from gail.discriminator import CnnDiscriminator, CnnDiscriminatorFlatAction +from src.discriminator import CnnDiscriminator, CnnDiscriminatorFlatAction +from src.policies import OptionsCnnPolicy + from imitation.algorithms import adversarial -import stable_baselines3 -import torch.utils.data -import numpy as np -from intersim.envs.intersimple import NRasterized -from intersim.collisions import state_to_polygon -import itertools -from torch.distributions import Categorical -import gym -import torch -import pickle +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 +from torch.distributions import Categorical +import numpy as np +import itertools +import gym +import pickle import tempfile import pathlib -from imitation.util import logger -from stable_baselines3.common.env_util import make_vec_env from tqdm import tqdm -model_name = 'gail_options_image' -env_settings = {'agent': 51, 'width': 36, 'height': 36, 'm_per_px': 2} +from intersim.envs.intersimple import NRasterized, NRasterizedRandomAgent, NRasterizedIncrementingAgent +from intersim.collisions import state_to_polygon 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) @@ -68,6 +51,13 @@ 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 @@ -104,13 +94,23 @@ 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, @@ -119,9 +119,18 @@ 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): @@ -131,11 +140,17 @@ 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()), @@ -145,6 +160,18 @@ 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, @@ -156,16 +183,29 @@ 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): @@ -326,8 +366,20 @@ def train_generator(env, generator, discriminator, num_samples): generator.train() -def train(expert_data, epochs=20, expert_batch_size=32, generator_steps=1024, discount=0.99): - env = NRasterized(**env_settings) +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") @@ -335,10 +387,10 @@ def train(expert_data, epochs=20, expert_batch_size=32, generator_steps=1024, di 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) + venv = make_vec_env(env_class, n_envs=1, env_kwargs=env_settings) discriminator = adversarial.GAIL( expert_data=expert_data, - expert_batch_size=expert_batch_size, + expert_batch_size=discrim_batch_size, discrim_kwargs={'discrim_net': CnnDiscriminatorFlatAction(venv)}, #discrim_kwargs={'discrim_net': CnnDiscriminator(venv)}, venv=venv, # unused @@ -360,7 +412,7 @@ def train(expert_data, epochs=20, expert_batch_size=32, generator_steps=1024, di ) for _ in tqdm(range(epochs)): - train_discriminator(LLOptions(env), generator, discriminator, num_samples=expert_batch_size) + train_discriminator(LLOptions(env), generator, discriminator, num_samples=discrim_batch_size) train_generator(HLOptions(env), generator, discriminator, num_samples=generator_steps) return generator @@ -368,18 +420,32 @@ def train(expert_data, epochs=20, expert_batch_size=32, generator_steps=1024, di # %% 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} + with open("data/NormalizedIntersimpleExpertMu.001_NRasterizedAgent51w36h36mppx2.pkl", "rb") as f: trajectories = pickle.load(f) transitions = rollout.flatten_trajectories(trajectories) - generator = train(transitions) + 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(NRasterized(**env_settings)) + env = RenderOptions(NRasterizedRandomAgent(**env_args)) for s in env.sample_ll(model): if s['dones']: diff --git a/src/policies/__init__.py b/src/policies/__init__.py index cff41ff..0679815 100644 --- a/src/policies/__init__.py +++ b/src/policies/__init__.py @@ -1 +1,2 @@ from src.policies.policy import IntersimPolicy, IntersimStateNet, IntersimStateActionNet, generate_transforms +from src.policies.options import OptionsCnnPolicy \ No newline at end of file diff --git a/src/policies/options.py b/src/policies/options.py new file mode 100644 index 0000000..9bf71c9 --- /dev/null +++ b/src/policies/options.py @@ -0,0 +1,59 @@ +import stable_baselines3 +from torch.distributions import Categorical + +class OptionsCnnPolicy(stable_baselines3.common.policies.ActorCriticCnnPolicy): + """ + Class for high-level options policy (generator) + """ + def __init__(self, observation_space, *args, **kwargs): + super().__init__(observation_space['obs'], *args, **kwargs) + + def _prior_distribution(self, s): + """ + Return prior distribution over high-level options (before masking) + Args: + s (torch.tensor): observation + Returns: + values (torch.tensor): values from critic + dist (torch.distributions): prior distribution over actions + """ + 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): + """ + Will mask invalid states before making action selections + Args: + obs: dict with keys: + obs (torch.tensor): (*,o) true observations + mask (torch.tensor): (*,m) mask over valid actions + Returns: + ch (torch.tensor): (*,a) sampled actions + values (torch.tensor): (*,) predicted value at observation + log_probs (torch.tensor): (*,) log probabilities of selected actions + """ + 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): + """ + Evaluate particular actions + Args: + obs: dict with keys: + obs (torch.tensor): (*,o) true observations + mask (torch.tensor): (*,m) masks over valid actions + ch (torch.tensor): (*,a) selected actions + Returns: + values (torch.tensor): (*,) predicted value at observation + log_probs (torch.tensor): (*,) log probabilities of selected actions + ent (torch.tensor): (*,) entropy of each distribution over actions + """ + 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 From 2d8928f2aedeafbc5ff4b727821d99d7138bd87e Mon Sep 17 00:00:00 2001 From: Arec Date: Thu, 21 Oct 2021 06:29:43 -0700 Subject: [PATCH 5/7] updating data processing scripts to output to the correct location --- generate_demos.sh | 4 +--- src/data/__init__.py | 2 +- src/data/data_utils.py | 2 +- src/data/expert.py | 2 +- 4 files changed, 4 insertions(+), 6 deletions(-) diff --git a/generate_demos.sh b/generate_demos.sh index 447068a..9900f5e 100755 --- a/generate_demos.sh +++ b/generate_demos.sh @@ -6,6 +6,4 @@ # expert_class:str='NormalizedIntersimpleExpert', # expert_args:dict={mu:0.001}): -cd ./src/data -python -m expert --locs='[DR_USA_Roundabout_FT]' -cd ../ \ No newline at end of file +python -m src.data.expert --locs='[DR_USA_Roundabout_FT]' \ No newline at end of file diff --git a/src/data/__init__.py b/src/data/__init__.py index 49d78bc..a134dbc 100644 --- a/src/data/__init__.py +++ b/src/data/__init__.py @@ -1 +1 @@ -from src.data.expert_trajectories import demonstrations, load_experts, process_experts \ No newline at end of file +from src.data.expert import demonstrations, load_experts, process_experts \ No newline at end of file diff --git a/src/data/data_utils.py b/src/data/data_utils.py index b14ec3a..0bcf9d4 100644 --- a/src/data/data_utils.py +++ b/src/data/data_utils.py @@ -1,7 +1,7 @@ import torch from torch.utils.data import Dataset import numpy as np -from src.expert_data import load_expert_data +from src.data.expert_data import load_expert_data import os opj = os.path.join diff --git a/src/data/expert.py b/src/data/expert.py index fb6b754..401531f 100644 --- a/src/data/expert.py +++ b/src/data/expert.py @@ -100,7 +100,7 @@ def save_video(env, expert): env.render() env.close() -def load_experts(expert_files=[]): +def load_experts(expert_files): """ Load expert trajectories from files and combine their transitions into a single RB From 06785236d43b972fbc3339ae1067a62134ba00f0 Mon Sep 17 00:00:00 2001 From: Arec Date: Thu, 21 Oct 2021 06:59:21 -0700 Subject: [PATCH 6/7] fixing expert data combiner and adding feasibility checking with all options to util.collisions --- src/data/expert.py | 7 +- src/util/__init__.py | 3 +- src/util/collisions.py | 155 +++++++++++++++++++++++++++++++++++++++++ 3 files changed, 161 insertions(+), 4 deletions(-) create mode 100644 src/util/collisions.py diff --git a/src/data/expert.py b/src/data/expert.py index 401531f..19ccee5 100644 --- a/src/data/expert.py +++ b/src/data/expert.py @@ -109,11 +109,12 @@ def load_experts(expert_files): Returns: transitions (list): list of combined expert episode transitions """ - transitions = [] + trajectories = [] for file in tqdm(expert_files): with open(file, "rb") as f: - trajectories = pickle.load(f) - transitions = transitions + rollout.flatten_trajectories(trajectories) + new_trajectories = pickle.load(f) + trajectories += new_trajectories + transitions = rollout.flatten_trajectories(trajectories) return transitions def demonstrations(expert='NormalizedIntersimpleExpert', env='NRasterizedIncrementingAgent', path=None, min_timesteps=None, min_episodes=None, video=False, env_args={}, policy_args={}): diff --git a/src/util/__init__.py b/src/util/__init__.py index 000365b..8515068 100644 --- a/src/util/__init__.py +++ b/src/util/__init__.py @@ -1 +1,2 @@ -from src.util.render_env import * \ No newline at end of file +from src.util.render_env import * +from src.util.collisions import feasible \ No newline at end of file diff --git a/src/util/collisions.py b/src/util/collisions.py new file mode 100644 index 0000000..8327857 --- /dev/null +++ b/src/util/collisions.py @@ -0,0 +1,155 @@ +import torch +import numpy as np +from intersim.collisions import state_to_polygon + +def feasible(env, plan, ch, method='exact'): + """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) + + # check_future_collisions_fast takes in B-list and outputs (B,) bool tensor + if method=='circle': + valid = check_future_collisions_fast(env, [full_plan]) + elif method=='ncircles': + valid = check_future_collisions_ncircles(env, [fullplan]) + elif method=='exact': + valid = check_future_collisions_exact(env, [full_plan]) + else: + raise NotImplementedError('Invalid collision-checking method') + return valid.item() + +def check_future_collisions_ncircles(env, actions, n_circles:int=2): + """Checks whether `env._agent` would collide with other agents assuming `actions` as input. + + Vehicles are (over-)approximated by multiple 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 + """ + assert n_circles >= 2 + 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) + centers = states[:, :, :, :2] + psi = states[:, :, :, 3] + lon = torch.stack([psi.cos(), psi.sin()],dim=-1) # (B, T, nv, 2) + + # offset between [-env._env.lengths+env._env.widths/2, env._env.lengths/2-env._env.widths/2] + back = (-env._env._lengths/2+env._env._widths/2).unsqueeze(-1) # (nv, 1) + length = (env._env._lengths-env._env._widths).unsqueeze(-1) # (nv, 1) + diff_d = back + length*(torch.arange(n_circles)/(n_circles-1)).unsqueeze(0) # (nv, n_circles) + assert diff_d.shape == (nv, n_circles) + + offsets = diff_d[None, None, :, :, None] * lon[:, :, :, None, :] + assert offsets.shape == (B, T, nv, n_circles, 2) + + expanded_centers=centers.unsqueeze(-2) + offsets #(B, T, nv, n_circles, 2) + assert expanded_centers.shape == (B, T, nv, n_circles, 2) + agent_centers = expanded_centers[:,:,env._agent:env._agent+1,:,:] #(B, T, 1, n_circles, 2) + ds = expanded_centers.reshape((B, T, nv*n_circles, 1, 2)) - agent_centers #(B, T, nv*nc,1, 2) - (B, T, 1, nc, 2) = (B, T, nv*nc, nc, 2) + + distance = (ds**2).sum(-1).sqrt().reshape((B, T, nv, n_circles, n_circles)) # (B, T, nv, nc, nc) + 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, n_circles, n_circles) + + radius = env._env._widths*np.sqrt(2) / 2 + min_distance = radius[env._agent] + radius + min_distance = min_distance[None, None, :, None, None] + assert min_distance.shape == (1, 1, nv, 1, 1) + + return (distance > min_distance).all(-1).all(-1).all(-1).all(-1) + +def check_future_collisions_circle(env, actions): + """Compute collision information for circular vehicle approximations + + Args: + env (gym.Env): current environment state + actions (list of torch.Tensor): list of B (T, nv, adims) T-length action profiles + Returns: + states (torch.Tensor): tensor of shape (B, T, nv, 5) of future states based on the action profiles + collision_tensor (torch.Tensor): tensor of shape (B, T, nv) of bools indicating which plan collides with which vehicles in which time frame + false: colliding, true: not colliding + """ + 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) + + collision_tensor = distance > min_distance + assert collision_tensor.shape == (B, T, nv) + return states, collision_tensor + +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 + """ + _, collision_tensor = check_future_collisions_circle(env, actions) + return collision_tensor.all(-1).all(-1) + +def check_future_collisions_exact(env, actions): + """ + Checks whether `env._agent` would collide with other agents assuming `actions` as input. + + 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 + """ + # First check with simple circle collision check + states, collision_tensor = check_future_collisions_circle(env, actions) + (B, T, nv, _) = states.shape + # For those that have colliding circles, check exactly + colliding_mask = ~collision_tensor + + ego_states = states[:, :, env._agent:env._agent+1, :].expand(states.shape) + assert ego_states.shape == states.shape + + # get dimensions + lengths = env._env._lengths.expand(states.shape[:3]) + widths = env._env._widths.expand(states.shape[:3]) + ego_lengths = lengths[:, :, env._agent:env._agent+1].expand(lengths.shape) + ego_widths = widths[:, :, env._agent:env._agent+1].expand(widths.shape) + assert lengths.shape == widths.shape == ego_lengths.shape == ego_widths.shape == (B, T, nv) + + # For every collision instance between ego and other vehicle, check whether rectangles intersect + exact_collisions = torch.zeros_like(collision_tensor[colliding_mask]) + for i, (ego_state, ego_length, ego_width, other_state, other_length, other_width) in enumerate(zip( + ego_states[colliding_mask], ego_lengths[colliding_mask], ego_widths[colliding_mask], + states[colliding_mask], lengths[colliding_mask], widths[colliding_mask] + )): + assert ego_state.shape == other_state.shape == (5,) + assert ego_length.shape == ego_width.shape == other_length.shape == other_width.shape == () + p_ego = state_to_polygon(ego_state, ego_length, ego_width) + p_other = state_to_polygon(other_state, other_length, other_width) + exact_collisions[i] = p_ego.intersects(p_other) + + collision_tensor[colliding_mask] = ~exact_collisions + return collision_tensor.all(-1).all(-1) + From ba79de58b86935b29613d83ee5768c9558654d57 Mon Sep 17 00:00:00 2001 From: Arec Date: Thu, 21 Oct 2021 07:01:12 -0700 Subject: [PATCH 7/7] moving feasibility checkers into src.util.collisions, and doing expert processing using the tools in src.data.expert --- .../etienne/intersimple/gail_options_image.py | 110 ++---------------- 1 file changed, 7 insertions(+), 103 deletions(-) diff --git a/scratch/etienne/intersimple/gail_options_image.py b/scratch/etienne/intersimple/gail_options_image.py index 150afe3..71d4741 100644 --- a/scratch/etienne/intersimple/gail_options_image.py +++ b/scratch/etienne/intersimple/gail_options_image.py @@ -1,6 +1,10 @@ # %% +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 @@ -11,7 +15,6 @@ from stable_baselines3.common.env_util import make_vec_env import torch import torch.utils.data -from torch.distributions import Categorical import numpy as np import itertools import gym @@ -21,7 +24,6 @@ import pathlib from tqdm import tqdm from intersim.envs.intersimple import NRasterized, NRasterizedRandomAgent, NRasterizedIncrementingAgent -from intersim.collisions import state_to_polygon ALL_OPTIONS = [(v,t) for v in [0,2,4,6,8] for t in [5, 10, 20]] # option 0 is safe fallback @@ -235,103 +237,6 @@ def generate_plan(env, i): assert len(plan) == t, "incorrect plan length" return plan -def check_future_collisions_circle(env, actions): - """Compute collision information for circular vehicle approximations - - Args: - env (gym.Env): current environment state - actions (list of torch.Tensor): list of B (T, nv, adims) T-length action profiles - Returns: - states (torch.Tensor): tensor of shape (B, T, nv, 5) of future states based on the action profiles - collision_tensor (torch.Tensor): tensor of shape (B, T, nv) of bools indicating which plan collides with which vehicles in which time frame - false: colliding, true: not colliding - """ - 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) - - collision_tensor = distance > min_distance - assert collision_tensor.shape == (B, T, nv) - return states, collision_tensor - -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 - """ - _, collision_tensor = check_future_collisions_circle(env, actions) - return collision_tensor.all(-1).all(-1) - -def check_future_collisions_exact(env, actions): - """ - Checks whether `env._agent` would collide with other agents assuming `actions` as input. - - 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 - """ - # First check with simple circle collision check - states, collision_tensor = check_future_collisions_circle(env, actions) - (B, T, nv, _) = states.shape - # For those that have colliding circles, check exactly - colliding_mask = ~collision_tensor - - ego_states = states[:, :, env._agent:env._agent+1, :].expand(states.shape) - assert ego_states.shape == states.shape - - # get dimensions - lengths = env._env._lengths.expand(states.shape[:3]) - widths = env._env._widths.expand(states.shape[:3]) - ego_lengths = lengths[:, :, env._agent:env._agent+1].expand(lengths.shape) - ego_widths = widths[:, :, env._agent:env._agent+1].expand(widths.shape) - assert lengths.shape == widths.shape == ego_lengths.shape == ego_widths.shape == (B, T, nv) - - # For every collision instance between ego and other vehicle, check whether rectangles intersect - exact_collisions = torch.zeros_like(collision_tensor[colliding_mask]) - for i, (ego_state, ego_length, ego_width, other_state, other_length, other_width) in enumerate(zip( - ego_states[colliding_mask], ego_lengths[colliding_mask], ego_widths[colliding_mask], - states[colliding_mask], lengths[colliding_mask], widths[colliding_mask] - )): - assert ego_state.shape == other_state.shape == (5,) - assert ego_length.shape == ego_width.shape == other_length.shape == other_width.shape == () - p_ego = state_to_polygon(ego_state, ego_length, ego_width) - p_other = state_to_polygon(other_state, other_length, other_width) - exact_collisions[i] = p_ego.intersects(p_other) - - collision_tensor[colliding_mask] = ~exact_collisions - return collision_tensor.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 - valid = check_future_collisions_exact(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), @@ -426,10 +331,9 @@ if __name__ == '__main__': #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) - with open("data/NormalizedIntersimpleExpertMu.001_NRasterizedAgent51w36h36mppx2.pkl", "rb") as f: - trajectories = pickle.load(f) - transitions = rollout.flatten_trajectories(trajectories) generator = train( transitions, env_class=env_class, @@ -438,7 +342,7 @@ if __name__ == '__main__': discrim_batch_size=32, generator_steps=2048, discount=0.99 - )) + ) generator.save(model_name)