diff --git a/eval_experiments.py b/eval_experiments.py index 109c5d5..bb552fd 100644 --- a/eval_experiments.py +++ b/eval_experiments.py @@ -6,7 +6,7 @@ import json activations = [torch.nn.Tanh, torch.nn.LeakyReLU] -def main(method:str='expert', folder:str=None, locations=[(0,0)], skip_running=False): +def main(method:str='expert', folder:str=None, locations=[(0,0)], skip_running=False, save_videos:bool=False, videos_folder:str='videos'): exclude_keys_from_policy_kwargs = {'learning_rate', 'learning_rate_decay', 'clip_ratio', 'iterations_per_epoch', 'option'} policy_kwargs = {} @@ -50,7 +50,8 @@ def main(method:str='expert', folder:str=None, locations=[(0,0)], skip_running=F policy_file=policy_file, policy_kwargs=policy_kwargs, env=env, - env_kwargs=env_kwargs) + env_kwargs=env_kwargs, + videos_folder=None if not save_videos else videos_folder) outfolder = os.path.dirname(outbase) else: locstr = 'loc_'+'_'.join([f'r{ro}t{tr}' for (ro,tr) in locations]) diff --git a/generate_videos.sh b/generate_videos.sh new file mode 100755 index 0000000..d3b94d5 --- /dev/null +++ b/generate_videos.sh @@ -0,0 +1,18 @@ +# can add --skip_running if you've already run the saved policies through the test environments and have appropriate +# metrics in the out folder. Doing so will generate average metrics quickly. + +# Experiment A +python -m eval_experiments +python -m eval_experiments --method idm --save_videos +python -m eval_experiments --method bc --folder='test_policies/bc/expA' --save_videos +python -m eval_experiments --method gail --folder='test_policies/gail/expA' --save_videos +python -m eval_experiments --method hail --folder='test_policies/hail/expA' --save_videos +python -m eval_experiments --method shail --folder='test_policies/shail/expA' --save_videos + +# Experiment B +python -m eval_experiments --locations='[(0,4)]' +python -m eval_experiments --method idm --locations='[(0,4)]' --save_videos +python -m eval_experiments --method bc --folder='test_policies/bc/expB' --locations='[(0,4)]' --save_videos +python -m eval_experiments --method gail --folder='test_policies/gail/expB' --locations='[(0,4)]' --save_videos +python -m eval_experiments --method hail --folder='test_policies/hail/expB' --locations='[(0,4)]' --save_videos +python -m eval_experiments --method shail --folder='test_policies/shail/expB' --locations='[(0,4)]' --save_videos diff --git a/src/__init__.py b/src/__init__.py index 08e852e..8b13789 100644 --- a/src/__init__.py +++ b/src/__init__.py @@ -1,2 +1 @@ -from src.data.expert_data import generate_expert_data, load_expert_data -from src.data.data_utils import InteractionDatasetSingleAgent \ No newline at end of file + diff --git a/src/eval_main.py b/src/eval_main.py index 37e78bd..d60f3fa 100644 --- a/src/eval_main.py +++ b/src/eval_main.py @@ -177,7 +177,8 @@ def evaluate_policy(locations:List[Tuple[int,int]], env_kwargs:dict, method: str, policy_file: str, - policy_kwargs:dict) -> List[Dict[str,list]]: + policy_kwargs:dict, + videos_folder: Optional[str] = None) -> List[Dict[str,list]]: """ Evaluate policy on an incrementing agent environment at all locations. Return metrics for that policy @@ -230,7 +231,13 @@ def evaluate_policy(locations:List[Tuple[int,int]], policy = load_policy(method, policy_file, policy_kwargs, eval_env) # run policy on environment - policy_metrics[i] = evaluator.evaluate(policy) + policy_videos_folder = None + if videos_folder is not None: + policy_videos_folder = os.path.join(videos_folder, method, f'loc{iround}', f'track{track}') + os.makedirs(policy_videos_folder, exist_ok=True) + policy_metrics[i] = evaluator.evaluate( + policy, videos_folder=policy_videos_folder + ) return policy_metrics @@ -364,7 +371,8 @@ def eval_main( policy_kwargs: dict={}, env: str='NRasterizedRouteIncrementingAgent', env_kwargs: dict={}, - seed: int=0): + seed: int=0, + videos_folder: Optional[str]=None): """ Test a particular model at different testing locations/tracks and compute average metrics over all files. @@ -410,7 +418,7 @@ def eval_main( else: # evaluate it on the given roundabouts - policy_metrics = evaluate_policy(locations, env, env_kwargs, method, policy_file, policy_kwargs) + policy_metrics = evaluate_policy(locations, env, env_kwargs, method, policy_file, policy_kwargs, videos_folder=videos_folder) smetrics = summary_metrics(policy_metrics) save_metrics(smetrics, outbase+'_summary.pkl') cmetrics = comparison_metrics(policy_metrics, expert_metrics, outbase=outbase) diff --git a/src/evaluation/evaluation.py b/src/evaluation/evaluation.py index abc6b34..b47f978 100644 --- a/src/evaluation/evaluation.py +++ b/src/evaluation/evaluation.py @@ -9,6 +9,7 @@ from tqdm import tqdm from src.util.wrappers import IntersimpleTimeLimit from src.options.envs import OptionsEnv from src.safe_options.options import SafeOptionsEnv +from src.evaluation.vec_env import CallbackWhenDoneVecEnv class IntersimpleEvaluation: """ @@ -80,7 +81,7 @@ class IntersimpleEvaluation: with open(filestr, 'wb') as f: pickle.dump(self._metrics, f) - def evaluate(self, policy, filestr: Optional[str] = None) -> Dict[str, list]: + def evaluate(self, policy, filestr: Optional[str] = None, videos_folder: Optional[str] = None) -> Dict[str, list]: """ Evaluate a policy on the incrementing agent evaluation environment @@ -88,6 +89,8 @@ class IntersimpleEvaluation: policy (BaseClass.BaseAlgorithm): policy in which policy.predict(observation)[0] returns an action filestr (str): path-like string to dump metrics to or None """ + self.videos_folder = videos_folder + self.reset() if self.use_pbar: self.pbar = tqdm(total=self.n_episodes) @@ -97,10 +100,11 @@ class IntersimpleEvaluation: evaluate_policy( policy, - self.env, + self.env if self.videos_folder is None else CallbackWhenDoneVecEnv([lambda: self.env], self.done_callback), n_eval_episodes=self.n_episodes, callback=self.evaluate_options_policy_callback if self.is_options_env else self.evaluate_policy_callback, - return_episode_rewards=False + return_episode_rewards=False, + render=self.videos_folder is not None, ) if self.use_pbar: self.pbar.close() @@ -145,6 +149,13 @@ class IntersimpleEvaluation: if done and self.use_pbar: self.pbar.update(1) + + def done_callback(self, info): + if self.is_options_env: + info = info['ll']['infos'][0] + agent = info['agent'] + filestr = os.path.join(self.videos_folder, f'agent{agent}') + self.env.close(filestr=filestr) def post_proc(self): """ diff --git a/src/evaluation/vec_env.py b/src/evaluation/vec_env.py new file mode 100644 index 0000000..047ef00 --- /dev/null +++ b/src/evaluation/vec_env.py @@ -0,0 +1,30 @@ +from stable_baselines3.common.vec_env import DummyVecEnv +from stable_baselines3.common.vec_env.base_vec_env import VecEnvStepReturn +from copy import deepcopy +import numpy as np + +class CallbackWhenDoneVecEnv(DummyVecEnv): + """DummyVecEnv that calls `done_callback` before resetting the wrapped environment.""" + + def __init__(self, env_fns, done_callback): + assert len(env_fns) == 1 # for now + super().__init__(env_fns) + self.done_callback = done_callback + + def step_wait(self) -> VecEnvStepReturn: + for env_idx in range(self.num_envs): + obs, self.buf_rews[env_idx], self.buf_dones[env_idx], self.buf_infos[env_idx] = self.envs[env_idx].step( + self.actions[env_idx] + ) + if self.buf_dones[env_idx]: + # save final observation where user can get it, then reset + self.buf_infos[env_idx]["terminal_observation"] = obs + + self.done_callback(deepcopy(self.buf_infos[env_idx])) + + obs = self.envs[env_idx].reset() + self._save_obs(env_idx, obs) + return (self._obs_from_buf(), np.copy(self.buf_rews), np.copy(self.buf_dones), deepcopy(self.buf_infos)) + + def render(self, mode='post'): + super().render(mode) diff --git a/src/options/envs.py b/src/options/envs.py index 6c58045..30b2e47 100644 --- a/src/options/envs.py +++ b/src/options/envs.py @@ -45,6 +45,7 @@ class OptionsEnv(Wrapper): self.options = options self.action_space = gym.spaces.Discrete(len(options)) self.max_plan_length = max(t for _, t in options) + self.render_mode = None def plan(self, option): target_v, t = option @@ -84,11 +85,14 @@ class OptionsEnv(Wrapper): n_steps = k + 1 return observations, actions, rewards, env_done, plan_done, infos, n_steps + + def render(self, mode='post'): + self.render_mode = mode - def step(self, action, render_mode=None): + def step(self, action): a = int(action) assert a == action - ll_obs, ll_actions, ll_rewards, ll_env_done, ll_plan_done, ll_infos, ll_steps = self.execute_plan(self.last_obs, self.options[a], render_mode) + ll_obs, ll_actions, ll_rewards, ll_env_done, ll_plan_done, ll_infos, ll_steps = self.execute_plan(self.last_obs, self.options[a], self.render_mode) hl_obs = ll_obs[ll_steps] hl_reward = (ll_rewards * ~ll_plan_done).sum().item() hl_done = ll_env_done[ll_steps-1].item() diff --git a/src/safe_options/options.py b/src/safe_options/options.py index 7e160af..55a0b70 100644 --- a/src/safe_options/options.py +++ b/src/safe_options/options.py @@ -225,8 +225,8 @@ class SafeOptionsEnv(OptionsEnv): } return obs - def step(self, action, render_mode=None): - obs, reward, done, info = super().step(action, render_mode) + def step(self, action): + obs, reward, done, info = super().step(action) obs = { 'observation': obs, 'safe_actions': self.safe_actions(), diff --git a/src/util/wrappers.py b/src/util/wrappers.py index d2d62ea..0d68da8 100644 --- a/src/util/wrappers.py +++ b/src/util/wrappers.py @@ -5,14 +5,23 @@ class Wrapper(gym.Wrapper): def __getattr__(self, name): return getattr(self.env, name) + def close(self, *args, **kwargs): + return self.env.close(*args, **kwargs) + class TransformObservation(gym.wrappers.TransformObservation): def __getattr__(self, name): return getattr(self.env, name) + def close(self, *args, **kwargs): + return self.env.close(*args, **kwargs) + class IntersimpleTimeLimit(gym.wrappers.TimeLimit): def __getattr__(self, name): return getattr(self.env, name) + def close(self, *args, **kwargs): + return self.env.close(*args, **kwargs) + class CollisionPenaltyWrapper(Wrapper): def __init__(self, env, collision_distance, collision_penalty, *args, **kwargs):