From da1fb112692748a2a264965a1fcb110a0298cfb1 Mon Sep 17 00:00:00 2001 From: Arec Date: Thu, 21 Oct 2021 05:21:06 -0700 Subject: [PATCH] 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