From 351b22fe01b8e54209790be6383762582f9d6539 Mon Sep 17 00:00:00 2001 From: Johannes Fischer Date: Thu, 28 Oct 2021 09:40:20 +0200 Subject: [PATCH] Optionally include infos in expert data --- scratch/etienne/intersimple/data/expert.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/scratch/etienne/intersimple/data/expert.py b/scratch/etienne/intersimple/data/expert.py index 56105af..e198502 100644 --- a/scratch/etienne/intersimple/data/expert.py +++ b/scratch/etienne/intersimple/data/expert.py @@ -1,4 +1,4 @@ -from intersim.envs.intersimple import Intersimple +from intersim.envs.intersimple import Intersimple, InfoFilter from stable_baselines3.common.policies import BasePolicy import gym import intersim.envs.intersimple @@ -117,6 +117,7 @@ def demonstrations(expert='NormalizedIntersimpleExpert', env='NRasterizedRandomA save_video(env, policy) path = path or (policy.__class__.__name__ + '_' + env.__class__.__name__ + '.pkl') + include_infos = isinstance(env, InfoFilter) rollout.rollout_and_save( path=path, @@ -125,7 +126,8 @@ def demonstrations(expert='NormalizedIntersimpleExpert', env='NRasterizedRandomA sample_until=rollout.make_sample_until( min_timesteps=min_timesteps, min_episodes=min_episodes, - ) + ), + exclude_infos=not include_infos, ) if __name__ == '__main__':