From 427a9e4f1b5a397b343de0a2adb6e8855ee63a93 Mon Sep 17 00:00:00 2001 From: Arec Date: Tue, 18 Jan 2022 17:53:56 -0800 Subject: [PATCH] adding rollout and trajectory flattening utilities. including old implementation in this commit --- generate_demos.sh | 3 +- src/data/expert.py | 56 ++++++++++-- src/util/rollout.py | 206 ++++++++++++++++++++++++++++++++++++++++++++ 3 files changed, 255 insertions(+), 10 deletions(-) create mode 100644 src/util/rollout.py diff --git a/generate_demos.sh b/generate_demos.sh index a7e2f31..7b06689 100755 --- a/generate_demos.sh +++ b/generate_demos.sh @@ -6,4 +6,5 @@ # expert_class:str='NRasterizedRouteIncrementingAgent', # expert_args:dict={mu:0.001}): -python -m src.data.expert --locs='[DR_USA_Roundabout_FT]' --tracks='[0]' \ No newline at end of file +# python -m src.data.expert --locs='[DR_USA_Roundabout_FT]' --tracks='[0]' +python -m src.data.expert --locs='[DR_USA_Roundabout_FT]' --tracks='[0]' --filename='expert_new.pkl' \ No newline at end of file diff --git a/src/data/expert.py b/src/data/expert.py index 0abd2c8..0b9ca90 100644 --- a/src/data/expert.py +++ b/src/data/expert.py @@ -5,13 +5,16 @@ 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 import numpy as np +import imitation.data.rollout as rollout +from imitation.data.wrappers import RolloutInfoWrapper + +from src.util.rollout import rollout_and_save, flatten_trajectories, make_sample_until + class IntersimExpert(BasePolicy): def __init__(self, intersim_env, mu=0, *args, **kwargs): @@ -149,12 +152,13 @@ def single_agent_expert(expert='NormalizedIntersimpleExpert', Expert = globals()[expert] env = Env(**env_args) policy = Expert(env, **policy_args) + # single_agent_demonstrations_old(env, policy, **kwargs) single_agent_demonstrations(env, policy, **kwargs) -def single_agent_demonstrations(env, policy, - path=None, min_timesteps=None, - min_episodes=None, video=False, - env_args={}, policy_args={}): +def single_agent_demonstrations_old(env, policy, + path=None, min_timesteps=None, + min_episodes=None, video=False, + env_args={}, policy_args={}): """Rollout and save expert demos. Usage: @@ -182,9 +186,8 @@ def single_agent_demonstrations(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, - ) + min_timesteps=min_timesteps, + min_episodes=min_episodes) rollout.rollout_and_save( path=path, @@ -194,6 +197,41 @@ def single_agent_demonstrations(env, policy, rng=NoShuffleRNG() ) +def single_agent_demonstrations(env, policy, + 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: + env (class): intersimple environment + policy (BasePolicy): intersimple policy + 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 + """ + + 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 = make_sample_until(min_timesteps=min_timesteps,min_episodes=min_episodes) + + rollout_and_save( + path=path, + policy=policy, + env=env, + sample_until=suntil + ) + def multi_agent_demonstrations(expert='IntersimExpert',path=None, env_args={}, policy_args={}): """ Run and save the `intersim' multiagent environment demonstration diff --git a/src/util/rollout.py b/src/util/rollout.py new file mode 100644 index 0000000..08e3522 --- /dev/null +++ b/src/util/rollout.py @@ -0,0 +1,206 @@ +# Borrowed heavily from https://github.com/HumanCompatibleAI/imitation/tree/master/src/imitation/data + +import pickle +import numpy as np +import logging +import os +import pathlib +from typing import Optional, List, Dict + +def generate_trajectories( + policy, + env, + sample_until, + rng: Optional[np.random.RandomState] = None, # np.random to shuffle +) -> List[dict]: + """Generate trajectory dictionaries from a policy and an environment. + + Args: + policy: a stable_baselines3 policy or algorithm trained on the gym environment + env: The environment to interact with. + sample_until: A function determining the termination condition. + It takes a sequence of trajectories, and returns a bool. + Most users will want to use one of `min_episodes` or `min_timesteps`. + rng: used for shuffling trajectories. + + Returns: + Sequence of trajectories, satisfying `sample_until`. + """ + trajectories = [] + while not sample_until(trajectories): + + # sample a trajectory + ob, done = env.reset(), False + ob_list, rew_list, act_list, info_list = [], [], [], [] + while not done: + act, _ = policy.predict(ob) + next_ob, rew, done, info = env.step(act) # ignore infos + ob_list.append(ob) + act_list.append(act) + rew_list.append(rew) + info_list.append(info) + ob = next_ob + ob_list.append(ob) + + traj = { + 'obs':np.stack(ob_list), + 'acts':np.stack(act_list), + 'rews':np.stack(rew_list), + 'infos': info_list, + 'terminal': True + } + trajectories.append(traj) + + # Shuffle trajectories + if rng: + rng.shuffle(trajectories) + + # Sanity checks. + for trajectory in trajectories: + n_steps = len(trajectory['acts']) + # extra 1 for the end + exp_obs = (n_steps + 1,) + env.observation_space.shape + real_obs = trajectory['obs'].shape + assert real_obs == exp_obs, f"expected shape {exp_obs}, got {real_obs}" + exp_act = (n_steps,) + env.action_space.shape + real_act = trajectory['acts'].shape + assert real_act == exp_act, f"expected shape {exp_act}, got {real_act}" + exp_rew = (n_steps,) + real_rew = trajectory['rews'].shape + assert real_rew == exp_rew, f"expected shape {exp_rew}, got {real_rew}" + + return trajectories + + +def flatten_trajectories(trajectory_list: list) -> list: + """ + Turn a list of trajectories into a (longer) list of transitions with appropriate fields + + Args: + trajectory_list (list): list of trajectory dicts with keys: + obs (np.ndarray): (T, *O) tensor of all observations in T-step trajectory + acts (np.ndarray): (T-1, *A) tensor of all actions in T-step trajectory + infos (list[dict]): (T-1)-length list of all information dictionaries + terminal (bool): True if the trajectory ends at the last step + rews (np.ndarray): (T-1, 1) tensor of rewards along trajectory + Returns: + transition_list (list): list of all transition dicts with keys: + obs (np.ndarray): (*O) tensor of single-step observation + acts (np.ndarray): (*A) tensor of single-step action + infos (dict): single-step information dictionary + next_obs (np.ndarray): (*O) tensor of next observation + rews (np.ndarray): (1) tensor of single-step reward + dones (bool): whether state is terminal + """ + transition_list = [] + for traj in trajectory_list: + T = traj['obs'].shape[0] + if traj['infos']: + infos = traj['infos'] + else: + infos = [{}] * T + for i in range(T-1): + transition_list.append({ + 'obs': traj['obs'][i], + 'acts': traj['acts'][i], + 'next_obs': traj['obs'][i+1], + 'rews': traj['rews'][i], + 'dones': (i==T-2), + 'infos': infos[i], + }) + return transition_list + + +def make_sample_until(min_timesteps: Optional[int]=None, min_episodes: Optional[int]=None): + """Returns a termination condition sampling for a number of timesteps and episodes. + + Args: + min_timesteps: Sampling will not stop until there are at least this many + timesteps. + min_episodes: Sampling will not stop until there are at least this many + episodes. + + Returns: + A termination condition which given a list of trajectories returns true if the condition is met. + + Raises: + ValueError if neither of n_timesteps and n_episodes are set, or if either are + non-positive. + """ + if min_timesteps is None and min_episodes is None: + raise ValueError( + "At least one of min_timesteps and min_episodes needs to be non-None" + ) + + conditions = [] + if min_timesteps is not None: + if min_timesteps <= 0: + raise ValueError( + f"min_timesteps={min_timesteps} if provided must be positive" + ) + + def timestep_cond(trajectories): + if len(trajectories) == 0: + return False + timesteps = sum(len(t['obs']) - 1 for t in trajectories) + return timesteps >= min_timesteps + conditions.append(timestep_cond) + + if min_episodes is not None: + if min_episodes <= 0: + raise ValueError( + f"min_episodes={min_episodes} if provided must be positive" + ) + conditions.append(lambda trajectories: len(trajectories) >= min_episodes) + + def sample_until(trajs: List[dict]) -> bool: + for cond in conditions: + if not cond(trajs): + return False + return True + + return sample_until + +def rollout_and_save( + path: str, + policy, + env, + sample_until, + *, + exclude_infos: bool = True, + **kwargs, +) -> None: + """Generate policy rollouts and save them to a pickled list of trajectories. + + The `.infos` field of each Trajectory is set to `None` to save space. + + Args: + path: Rollouts are saved to this path. + policy: a stable_baselines3 policy or algorithm trained on the gym environment + env: The environment to interact with. + sample_until: End condition for rollout sampling. + exclude_infos: If True, then exclude `infos` from pickle by setting + this field to None. Excluding `infos` can save a lot of space during + pickles. + **kwargs: Passed through to `generate_trajectories`. + """ + trajs = generate_trajectories(policy, env, sample_until, **kwargs) + if exclude_infos: + [traj.update(infos=None) for traj in trajs] + save(path, trajs) + +def save(path: str, trajectories: List[dict]) -> None: + """Save a sequence of Trajectories to disk. + + Args: + path: Trajectories are saved to this path. + trajectories: The trajectories to save. + """ + p = pathlib.Path(path) + p.parent.mkdir(parents=True, exist_ok=True) + tmp_path = f"{path}.tmp" + with open(tmp_path, "wb") as f: + pickle.dump(trajectories, f) + # Ensure atomic write + os.replace(tmp_path, path) + logging.info(f"Dumped demonstrations to {path}.") \ No newline at end of file