huangfu/InteractionImitation
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

adding rollout and trajectory flattening utilities. including old implementation in this commit

Browse Source
This commit is contained in:
Arec
2022-01-18 17:53:56 -08:00
parent 2c1dc6ca33
commit 427a9e4f1b
3 changed files with 255 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
generate_demos.sh
View File

@@ -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]'
# 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'

48
src/data/expert.py
View File

@@ -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,9 +152,10 @@ 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,
def single_agent_demonstrations_old(env, policy,
path=None, min_timesteps=None,
min_episodes=None, video=False,
env_args={}, policy_args={}):
@@ -183,8 +187,7 @@ 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_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 <flags>
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

206
src/util/rollout.py Normal file
View File

@@ -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}.")