From d34fa5774d42a7959bf2c0ef49752f4a8cbf402e Mon Sep 17 00:00:00 2001 From: Arec Date: Mon, 29 Nov 2021 14:09:57 -0800 Subject: [PATCH] adding functions to save joint expert states and actions for repeated use in metrics, adding option to flatten loaded trajectories, adding class to not shuffle trajectories when saving experts to make sure it lines up with the joint states. checked that it does --- scratch/arec/intersimple/commands.txt | 23 ++++++++- src/data/__init__.py | 2 +- src/data/expert.py | 73 +++++++++++++++++++++++---- 3 files changed, 87 insertions(+), 11 deletions(-) diff --git a/scratch/arec/intersimple/commands.txt b/scratch/arec/intersimple/commands.txt index b706521..aa91797 100644 --- a/scratch/arec/intersimple/commands.txt +++ b/scratch/arec/intersimple/commands.txt @@ -1 +1,22 @@ -python -m render_options --model_name='gail_options_image_mid_wcollision' --env='NRasterizedRoute' --options=True --width=36 --height=36 --m_per_px=2 --agent=50 --stop_on_collision=False \ No newline at end of file +python -m render_options --model_name='gail_options_image_mid_wcollision' --env='NRasterizedRoute' --options=True --width=36 --height=36 --m_per_px=2 --agent=50 --stop_on_collision=False + +import torch, os +from src.data import load_experts +folder = 'expert_data/DR_USA_Roundabout_FT/track0000' +single_agent = os.path.join(folder, 'expert.pkl') +multi_agent = os.path.join(folder,'joint_expert_states.pt') +multi_agent_actions = os.path.join(folder,'joint_expert_actions.pt') +demonstrations = load_experts([single_agent], flatten=False) +demonstrations[0].__dict__.keys() +len(demonstrations[0].obs) +single_agent_lengths = [len(demonstration.obs) for demonstration in demonstrations] +states = torch.load(multi_agent) +actions = torch.load(multi_agent_actions) +multi_agent_lengths = [sum(~torch.isnan(states[:,i,0])).item() for i in range(states.shape[1])] + +single_agent_actions = [demonstration.acts for demonstration in demonstrations] +multi_agent_actions = [actions[~torch.isnan(actions[:,i,0])] for i in range(actions.shape[1])] + +import pickle +with open(single_agent, "rb") as f: + new_trajectories = pickle.load(f) \ No newline at end of file diff --git a/src/data/__init__.py b/src/data/__init__.py index a134dbc..e2f723c 100644 --- a/src/data/__init__.py +++ b/src/data/__init__.py @@ -1 +1 @@ -from src.data.expert import demonstrations, load_experts, process_experts \ No newline at end of file +from src.data.expert import single_agent_demonstrations, multi_agent_demonstrations, NoShuffleRNG, load_experts, process_experts \ No newline at end of file diff --git a/src/data/expert.py b/src/data/expert.py index 6b35d83..e06649b 100644 --- a/src/data/expert.py +++ b/src/data/expert.py @@ -1,3 +1,4 @@ +import intersim from intersim.envs.intersimple import Intersimple from stable_baselines3.common.policies import BasePolicy import gym @@ -9,6 +10,7 @@ from stable_baselines3.common.vec_env.dummy_vec_env import DummyVecEnv from imitation.data.wrappers import RolloutInfoWrapper import copy import os +import numpy as np class IntersimExpert(BasePolicy): @@ -100,24 +102,40 @@ def save_video(env, expert): env.render() env.close() -def load_experts(expert_files): +class NoShuffleRNG(np.random.RandomState): + """ + A np.random.RandomState rng that doesn't shuffle inputs (for imitation.rollout) + """ + def __init__(self): + super().__init__() + + def shuffle(self, x): + return x + +def load_experts(expert_files, flatten = True): """ Load expert trajectories from files and combine their transitions into a single RB Args: expert_files (list): list of expert file strings + flatten (bool): whether to flatten trajectory info Returns: transitions (list): list of combined expert episode transitions """ - trajectories = [] + transitions = [] for file in tqdm(expert_files): with open(file, "rb") as f: new_trajectories = pickle.load(f) - trajectories += new_trajectories - transitions = rollout.flatten_trajectories(trajectories) + transitions += new_trajectories + if flatten: + transitions = rollout.flatten_trajectories(transitions) return transitions -def demonstrations(expert='NormalizedIntersimpleExpert', env='NRasterizedRouteIncrementingAgent', path=None, min_timesteps=None, min_episodes=None, video=False, env_args={}, policy_args={}): +def single_agent_demonstrations(expert='NormalizedIntersimpleExpert', + env='NRasterizedRouteIncrementingAgent', + path=None, min_timesteps=None, + min_episodes=None, video=False, + env_args={}, policy_args={}): """Rollout and save expert demos. Usage: @@ -154,13 +172,39 @@ def demonstrations(expert='NormalizedIntersimpleExpert', env='NRasterizedRouteIn min_timesteps=min_timesteps, min_episodes=min_episodes, ) + rollout.rollout_and_save( path=path, policy=venv_policy, venv=venv, - sample_until=suntil + sample_until=suntil, + rng=NoShuffleRNG() ) +def multi_agent_demonstrations(expert='IntersimExpert',path=None, env_args={}, policy_args={}): + """ + Run and save the `intersim' multiagent environment demonstration + + Args: + expert (class): class of multi-agent expert + path (str): path to store output data + env_args (dict): dictionary of kwargs when instantiating environment class + policy_args (dict): dictionary of kwargs when instantiating Expert policy + """ + if path is None: + raise('No path specified') + + env = gym.make('intersim:intersim-v0',**env_args) + Expert = globals()[expert] + policy = Expert(env, **policy_args) + + s, done = env.reset(), False + env.render(mode='file') + while not done: + _,_,done,_ = env.step(policy.predict()[0]) + env.render(mode='file') + env.close(filestr=path) + def process_experts(filename:str='expert.pkl', locs:list=None, tracks:list=None, @@ -190,22 +234,33 @@ def process_experts(filename:str='expert.pkl', iloc = intersim.LOCATIONS.index(loc) it_env_args = copy.deepcopy(env_args) - it_env_args.update({ + env_loc_args = { 'loc':iloc, 'track':track, - }) + } + it_env_args.update(env_loc_args) 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( + # Multi-Agent demonstrations + it_ma_path = os.path.join(out_folder,'joint_expert') + multi_agent_demonstrations( + expert='IntersimExpert', + path=it_ma_path, + env_args=env_loc_args, + policy_args=expert_args + ) + # Single-Agent POV Demonstrations + single_agent_demonstrations( expert=expert_class, env=env_class, path=it_path, env_args=it_env_args, policy_args=expert_args, ) + pbar.update(1) pbar.close()