conditional StreamWrapper

Signed-off-by: Mathieu D <[email protected]>

baselines/baselines_utils.py

@@ -35,7 +35,7 @@ def load_or_create_model(model_to_load_path, env_config, total_timesteps, num_cp
     return model
 
 
-def make_env(rank, env_conf, seed=0):
+def make_env(rank, env_config, seed=0):
     """
     Utility function for multiprocessed env.
     :param env_id: (str) the environment ID
@@ -44,8 +44,18 @@ def make_env(rank, env_conf, seed=0):
     :param rank: (int) index of the subprocess
     """
     def _init():
-        env = RedGymEnv(env_conf)
+        env = RedGymEnv(env_config)
         env.reset(seed=(seed + rank))
+        if env_config['stream'] is True:
+            env = StreamWrapper(
+                env,
+                stream_metadata = { # All of this is part is optional
+                    "user": "MATHIEU", # choose your own username
+                    "env_id": env_config['instance_id'], # environment identifier
+                    "color": "#d900ff", # choose your color :)
+                    "extra": "", # any extra text you put here will be displayed
+                }
+            )
         return env
     set_random_seed(seed)
     return _init